| 1 | # |
| 2 | # Attempt to find the python package that uses the same python executable as |
| 3 | # `EXECUTABLE` and is one of the `SUPPORTED_VERSIONS`. |
| 4 | # |
| 5 | macro (find_python_from_executable EXECUTABLE SUPPORTED_VERSIONS) |
| 6 | file(REAL_PATH ${EXECUTABLE} EXECUTABLE) |
| 7 | set(Python_EXECUTABLE ${EXECUTABLE}) |
| 8 | find_package(Python COMPONENTS Interpreter Development.Module Development.SABIModule) |
| 9 | if (NOT Python_FOUND) |
| 10 | message(FATAL_ERROR "Unable to find python matching: ${EXECUTABLE}.") |
| 11 | endif() |
| 12 | set(_VER "${Python_VERSION_MAJOR}.${Python_VERSION_MINOR}") |
| 13 | set(_SUPPORTED_VERSIONS_LIST ${SUPPORTED_VERSIONS} ${ARGN}) |
| 14 | if (NOT _VER IN_LIST _SUPPORTED_VERSIONS_LIST) |
| 15 | message(FATAL_ERROR |
| 16 | "Python version (${_VER}) is not one of the supported versions: " |
| 17 | "${_SUPPORTED_VERSIONS_LIST}.") |
| 18 | endif() |
| 19 | message(STATUS "Found python matching: ${EXECUTABLE}.") |
| 20 | endmacro() |
| 21 | |
| 22 | # |
| 23 | # Run `EXPR` in python. The standard output of python is stored in `OUT` and |
| 24 | # has trailing whitespace stripped. If an error is encountered when running |
| 25 | # python, a fatal message `ERR_MSG` is issued. |
| 26 | # |
| 27 | function (run_python OUT EXPR ERR_MSG) |
| 28 | execute_process( |
| 29 | COMMAND |
| 30 | "${Python_EXECUTABLE}" "-c" "${EXPR}" |
| 31 | OUTPUT_VARIABLE PYTHON_OUT |
| 32 | RESULT_VARIABLE PYTHON_ERROR_CODE |
| 33 | ERROR_VARIABLE PYTHON_STDERR |
| 34 | OUTPUT_STRIP_TRAILING_WHITESPACE) |
| 35 | |
| 36 | if(NOT PYTHON_ERROR_CODE EQUAL 0) |
| 37 | message(FATAL_ERROR "${ERR_MSG}: ${PYTHON_STDERR}") |
| 38 | endif() |
| 39 | set(${OUT} ${PYTHON_OUT} PARENT_SCOPE) |
| 40 | endfunction() |
| 41 | |
| 42 | # Run `EXPR` in python after importing `PKG`. Use the result of this to extend |
| 43 | # `CMAKE_PREFIX_PATH` so the torch cmake configuration can be imported. |
| 44 | macro (append_cmake_prefix_path PKG EXPR) |
| 45 | run_python(_PREFIX_PATH |
| 46 | "import ${PKG}; print(${EXPR})" "Failed to locate ${PKG} path") |
| 47 | list(APPEND CMAKE_PREFIX_PATH ${_PREFIX_PATH}) |
| 48 | endmacro() |
| 49 | |
| 50 | # Resolve hipified output paths for `NAME` into `OUT_SRCS` and register the |
| 51 | # `.cu` sources with the shared `hipify_all` target. Per-extension hipify |
| 52 | # targets are unsafe to run in parallel against a shared csrc/ output dir, so |
| 53 | # accumulation here is paired with a single finalize step. |
| 54 | function (hipify_sources_target OUT_SRCS NAME ORIG_SRCS) |
| 55 | if (TARGET hipify_all) |
| 56 | message(FATAL_ERROR |
| 57 | "hipify_sources_target(${NAME}) called after vllm_finalize_hipify_target. " |
| 58 | "Add the new HIP extension before the finalizer call in CMakeLists.txt.") |
| 59 | endif() |
| 60 | |
| 61 | # |
| 62 | # Split into C++ and non-C++ (i.e. CUDA) sources. |
| 63 | # |
| 64 | set(SRCS ${ORIG_SRCS}) |
| 65 | set(CXX_SRCS ${ORIG_SRCS}) |
| 66 | list(FILTER SRCS EXCLUDE REGEX "\.(cc)|(cpp)|(hip)$") |
| 67 | list(FILTER CXX_SRCS INCLUDE REGEX "\.(cc)|(cpp)|(hip)$") |
| 68 | |
| 69 | # |
| 70 | # Generate ROCm/HIP source file names from CUDA file names. |
| 71 | # Since HIP files are generated code, they will appear in the build area |
| 72 | # `CMAKE_CURRENT_BINARY_DIR` directory rather than the original csrc dir. |
| 73 | # |
| 74 | set(HIP_SRCS) |
| 75 | foreach (SRC ${SRCS}) |
| 76 | string(REGEX REPLACE "\.cu$" "\.hip" SRC ${SRC}) |
| 77 | string(REGEX REPLACE "cuda" "hip" SRC ${SRC}) |
| 78 | list(APPEND HIP_SRCS "${CMAKE_CURRENT_BINARY_DIR}/${SRC}") |
| 79 | endforeach() |
| 80 | |
| 81 | set_property(GLOBAL APPEND PROPERTY VLLM_HIPIFY_ALL_SRCS ${SRCS}) |
| 82 | set_property(GLOBAL APPEND PROPERTY VLLM_HIPIFY_ALL_BYPRODUCTS ${HIP_SRCS}) |
| 83 | |
| 84 | # Chain hipify targets so they run sequentially. Parallel hipify |
| 85 | # invocations race on shutil.copytree, overwriting .hip files |
| 86 | # produced by another target back to .cu originals. |
| 87 | if (DEFINED _VLLM_LAST_HIPIFY_TARGET) |
| 88 | add_dependencies(hipify${NAME} ${_VLLM_LAST_HIPIFY_TARGET}) |
| 89 | endif() |
| 90 | set(_VLLM_LAST_HIPIFY_TARGET "hipify${NAME}" PARENT_SCOPE) |
| 91 | |
| 92 | # Swap out original extension sources with hipified sources. |
| 93 | list(APPEND HIP_SRCS ${CXX_SRCS}) |
| 94 | set(${OUT_SRCS} ${HIP_SRCS} PARENT_SCOPE) |
| 95 | endfunction() |
| 96 | |
| 97 | # Define the single shared `hipify_all` custom target that runs hipify once |
| 98 | # on the union of every HIP extension's sources. Call after the last HIP |
| 99 | # `define_extension_target`. |
| 100 | function (vllm_finalize_hipify_target) |
| 101 | if (TARGET hipify_all) |
| 102 | return() |
| 103 | endif() |
| 104 | |
| 105 | get_property(ALL_SRCS GLOBAL PROPERTY VLLM_HIPIFY_ALL_SRCS) |
| 106 | get_property(ALL_BYPRODUCTS GLOBAL PROPERTY VLLM_HIPIFY_ALL_BYPRODUCTS) |
| 107 | |
| 108 | if (NOT ALL_SRCS) |
| 109 | return() |
| 110 | endif() |
| 111 | |
| 112 | list(REMOVE_DUPLICATES ALL_SRCS) |
| 113 | list(REMOVE_DUPLICATES ALL_BYPRODUCTS) |
| 114 | |
| 115 | set(CSRC_BUILD_DIR ${CMAKE_CURRENT_BINARY_DIR}/csrc) |
| 116 | add_custom_target( |
| 117 | hipify_all |
| 118 | COMMAND ${Python_EXECUTABLE} ${CMAKE_SOURCE_DIR}/cmake/hipify.py -p ${CMAKE_SOURCE_DIR}/csrc -o ${CSRC_BUILD_DIR} ${ALL_SRCS} |
| 119 | DEPENDS ${CMAKE_SOURCE_DIR}/cmake/hipify.py ${ALL_SRCS} |
| 120 | BYPRODUCTS ${ALL_BYPRODUCTS} |
| 121 | COMMENT "Running hipify on all extension source files.") |
| 122 | endfunction() |
| 123 | |
| 124 | # |
| 125 | # Get additional GPU compiler flags from torch. |
| 126 | # |
| 127 | function (get_torch_gpu_compiler_flags OUT_GPU_FLAGS GPU_LANG) |
| 128 | if (${GPU_LANG} STREQUAL "CUDA") |
| 129 | # |
| 130 | # Get common NVCC flags from torch. |
| 131 | # |
| 132 | run_python(GPU_FLAGS |
| 133 | "from torch.utils.cpp_extension import COMMON_NVCC_FLAGS; print(';'.join(COMMON_NVCC_FLAGS))" |
| 134 | "Failed to determine torch nvcc compiler flags") |
| 135 | |
| 136 | if (CUDA_VERSION VERSION_GREATER_EQUAL 11.8) |
| 137 | list(APPEND GPU_FLAGS "-DENABLE_FP8") |
| 138 | endif() |
| 139 | if (CUDA_VERSION VERSION_GREATER_EQUAL 12.0) |
| 140 | list(REMOVE_ITEM GPU_FLAGS |
| 141 | "-D__CUDA_NO_HALF_OPERATORS__" |
| 142 | "-D__CUDA_NO_HALF_CONVERSIONS__" |
| 143 | "-D__CUDA_NO_BFLOAT16_CONVERSIONS__" |
| 144 | "-D__CUDA_NO_HALF2_OPERATORS__") |
| 145 | endif() |
| 146 | |
| 147 | elseif(${GPU_LANG} STREQUAL "HIP") |
| 148 | # |
| 149 | # Get common HIP/HIPCC flags from torch. |
| 150 | # |
| 151 | run_python(GPU_FLAGS |
| 152 | "import torch.utils.cpp_extension as t; print(';'.join(t.COMMON_HIP_FLAGS + t.COMMON_HIPCC_FLAGS))" |
| 153 | "Failed to determine torch nvcc compiler flags") |
| 154 | |
| 155 | list(APPEND GPU_FLAGS |
| 156 | "-DUSE_ROCM" |
| 157 | "-DENABLE_FP8" |
| 158 | "-U__HIP_NO_HALF_CONVERSIONS__" |
| 159 | "-U__HIP_NO_HALF_OPERATORS__" |
| 160 | "-Werror=unused-variable" |
| 161 | "-fno-gpu-rdc") |
| 162 | |
| 163 | endif() |
| 164 | set(${OUT_GPU_FLAGS} ${GPU_FLAGS} PARENT_SCOPE) |
| 165 | endfunction() |
| 166 | |
| 167 | # Find libgomp that gets shipped with PyTorch wheel and create a shim dir with: |
| 168 | # libgomp.so -> libgomp-<hash>.so... |
| 169 | # libgomp.so.1 -> libgomp-<hash>.so... |
| 170 | # OUTPUT: TORCH_GOMP_SHIM_DIR ("" if not found) |
| 171 | function(vllm_prepare_torch_gomp_shim TORCH_GOMP_SHIM_DIR) |
| 172 | set(${TORCH_GOMP_SHIM_DIR} "" PARENT_SCOPE) |
| 173 | |
| 174 | # Use run_python to locate vendored libgomp; never throw on failure. |
| 175 | run_python(_VLLM_TORCH_GOMP_PATH |
| 176 | " |
| 177 | import os, glob |
| 178 | import torch |
| 179 | torch_pkg = os.path.dirname(torch.__file__) |
| 180 | site_root = os.path.dirname(torch_pkg) |
| 181 | |
| 182 | # Search both torch.libs and torch/lib |
| 183 | roots = [os.path.join(site_root, 'torch.libs'), os.path.join(torch_pkg, 'lib')] |
| 184 | candidates = [] |
| 185 | for root in roots: |
| 186 | if not os.path.isdir(root): |
| 187 | continue |
| 188 | candidates.extend(glob.glob(os.path.join(root, 'libgomp*.so*'))) |
| 189 | |
| 190 | print(candidates[0] if candidates else '') |
| 191 | " |
| 192 | "failed to probe for libgomp") |
| 193 | |
| 194 | if(_VLLM_TORCH_GOMP_PATH STREQUAL "" OR NOT EXISTS "${_VLLM_TORCH_GOMP_PATH}") |
| 195 | return() |
| 196 | endif() |
| 197 | |
| 198 | # Create shim under the build tree |
| 199 | set(_shim "${CMAKE_BINARY_DIR}/gomp_shim") |
| 200 | file(MAKE_DIRECTORY "${_shim}") |
| 201 | |
| 202 | execute_process(COMMAND ${CMAKE_COMMAND} -E rm -f "${_shim}/libgomp.so") |
| 203 | execute_process(COMMAND ${CMAKE_COMMAND} -E rm -f "${_shim}/libgomp.so.1") |
| 204 | execute_process(COMMAND ${CMAKE_COMMAND} -E create_symlink "${_VLLM_TORCH_GOMP_PATH}" "${_shim}/libgomp.so") |
| 205 | execute_process(COMMAND ${CMAKE_COMMAND} -E create_symlink "${_VLLM_TORCH_GOMP_PATH}" "${_shim}/libgomp.so.1") |
| 206 | |
| 207 | set(${TORCH_GOMP_SHIM_DIR} "${_shim}" PARENT_SCOPE) |
| 208 | endfunction() |
| 209 | |
| 210 | # Macro for converting a `gencode` version number to a cmake version number. |
| 211 | # Preserves architecture-specific suffixes (a/f) needed for correct |
| 212 | # __CUDA_ARCH_FAMILY_SPECIFIC__ definition. E.g. "121a" -> "12.1a". |
| 213 | macro(string_to_ver OUT_VER IN_STR) |
| 214 | string(REGEX REPLACE "\([0-9]+\)\([0-9][af]?\)" "\\1.\\2" ${OUT_VER} ${IN_STR}) |
| 215 | endmacro() |
| 216 | |
| 217 | # |
| 218 | # Clear all `-gencode` flags from `CMAKE_CUDA_FLAGS` and store them in |
| 219 | # `CUDA_ARCH_FLAGS`. |
| 220 | # |
| 221 | # Example: |
| 222 | # CMAKE_CUDA_FLAGS="-Wall -gencode arch=compute_70,code=sm_70 -gencode arch=compute_75,code=sm_75" |
| 223 | # clear_cuda_gencode_flags(CUDA_ARCH_FLAGS) |
| 224 | # CUDA_ARCH_FLAGS="-gencode arch=compute_70,code=sm_70;-gencode arch=compute_75,code=sm_75" |
| 225 | # CMAKE_CUDA_FLAGS="-Wall" |
| 226 | # |
| 227 | macro(clear_cuda_gencode_flags CUDA_ARCH_FLAGS) |
| 228 | # Extract all `-gencode` flags from `CMAKE_CUDA_FLAGS` |
| 229 | string(REGEX MATCHALL "-gencode arch=[^ ]+" CUDA_ARCH_FLAGS |
| 230 | ${CMAKE_CUDA_FLAGS}) |
| 231 | |
| 232 | # Remove all `-gencode` flags from `CMAKE_CUDA_FLAGS` since they will be modified |
| 233 | # and passed back via the `CUDA_ARCHITECTURES` property. |
| 234 | string(REGEX REPLACE "-gencode arch=[^ ]+ *" "" CMAKE_CUDA_FLAGS |
| 235 | ${CMAKE_CUDA_FLAGS}) |
| 236 | endmacro() |
| 237 | |
| 238 | # |
| 239 | # Warn when a caller requested PTX code generation through global CUDA arch |
| 240 | # flags. vLLM removes those flags and reapplies per-source gencode flags, so the |
| 241 | # user's global PTX request will not be preserved. |
| 242 | # |
| 243 | function(warn_if_ptx_arch_requested CUDA_ARCH_FLAGS) |
| 244 | foreach(_ARCH_FLAG ${CUDA_ARCH_FLAGS}) |
| 245 | if(_ARCH_FLAG MATCHES "code=.*compute_[0-9]+[af]?") |
| 246 | message(WARNING |
| 247 | "PTX code generation requested in CUDA architecture flags " |
| 248 | "(${_ARCH_FLAG}), but vLLM does not preserve global PTX requests " |
| 249 | "when normalizing per-source CUDA architectures. Remove '+PTX' from " |
| 250 | "TORCH_CUDA_ARCH_LIST or rely on vLLM's built-in per-kernel PTX " |
| 251 | "selection.") |
| 252 | return() |
| 253 | endif() |
| 254 | endforeach() |
| 255 | endfunction() |
| 256 | |
| 257 | |
| 258 | # |
| 259 | # Extract unique CUDA architectures from a list of compute capabilities codes in |
| 260 | # the form `<major><minor>[<letter>]`, convert them to the form sort |
| 261 | # `<major>.<minor>`, dedupes them and then sorts them in ascending order and |
| 262 | # stores them in `OUT_ARCHES`. |
| 263 | # |
| 264 | # Prefer `code=sm_*`; fall back to `arch=compute_*` for PTX-only flags. |
| 265 | # This handles mismatches such as `arch=compute_20,code=sm_121`. |
| 266 | function(extract_unique_cuda_archs_ascending OUT_ARCHES CUDA_ARCH_FLAGS) |
| 267 | set(_CUDA_ARCHES) |
| 268 | foreach(_ARCH ${CUDA_ARCH_FLAGS}) |
| 269 | string(REGEX MATCH "code=sm_\([0-9]+[af]?\)" _COMPUTE ${_ARCH}) |
| 270 | if (NOT _COMPUTE) |
| 271 | string(REGEX MATCH "arch=compute_\([0-9]+[af]?\)" _COMPUTE ${_ARCH}) |
| 272 | endif() |
| 273 | if (_COMPUTE) |
| 274 | set(_COMPUTE ${CMAKE_MATCH_1}) |
| 275 | endif() |
| 276 | |
| 277 | string_to_ver(_COMPUTE_VER ${_COMPUTE}) |
| 278 | list(APPEND _CUDA_ARCHES ${_COMPUTE_VER}) |
| 279 | endforeach() |
| 280 | |
| 281 | list(REMOVE_DUPLICATES _CUDA_ARCHES) |
| 282 | list(SORT _CUDA_ARCHES COMPARE NATURAL ORDER ASCENDING) |
| 283 | set(${OUT_ARCHES} ${_CUDA_ARCHES} PARENT_SCOPE) |
| 284 | endfunction() |
| 285 | |
| 286 | # |
| 287 | # For a specific file set the `-gencode` flag in compile options conditionally |
| 288 | # for the CUDA language. |
| 289 | # |
| 290 | # Example: |
| 291 | # set_gencode_flag_for_srcs( |
| 292 | # SRCS "foo.cu" |
| 293 | # ARCH "compute_75" |
| 294 | # CODE "sm_75") |
| 295 | # adds: "-gencode arch=compute_75,code=sm_75" to the compile options for |
| 296 | # `foo.cu` (only for the CUDA language). |
| 297 | # |
| 298 | macro(set_gencode_flag_for_srcs) |
| 299 | set(options) |
| 300 | set(oneValueArgs ARCH CODE) |
| 301 | set(multiValueArgs SRCS) |
| 302 | cmake_parse_arguments(arg "${options}" "${oneValueArgs}" |
| 303 | "${multiValueArgs}" ${ARGN} ) |
| 304 | set(_FLAG -gencode arch=${arg_ARCH},code=${arg_CODE}) |
| 305 | set_property( |
| 306 | SOURCE ${arg_SRCS} |
| 307 | APPEND PROPERTY |
| 308 | COMPILE_OPTIONS "$<$<COMPILE_LANGUAGE:CUDA>:${_FLAG}>" |
| 309 | ) |
| 310 | |
| 311 | message(DEBUG "Setting gencode flag for ${arg_SRCS}: ${_FLAG}") |
| 312 | endmacro(set_gencode_flag_for_srcs) |
| 313 | |
| 314 | # |
| 315 | # For a list of source files set the `-gencode` flags in the files specific |
| 316 | # compile options (specifically for the CUDA language). |
| 317 | # |
| 318 | # arguments are: |
| 319 | # SRCS: list of source files |
| 320 | # CUDA_ARCHS: list of CUDA architectures in the form `<major>.<minor>[letter]` |
| 321 | # BUILD_PTX_FOR_ARCH: if set to true, then the PTX code will be built |
| 322 | # for architecture `BUILD_PTX_FOR_ARCH` if there is a CUDA_ARCH in CUDA_ARCHS |
| 323 | # that is larger than BUILD_PTX_FOR_ARCH. |
| 324 | # |
| 325 | macro(set_gencode_flags_for_srcs) |
| 326 | set(options) |
| 327 | set(oneValueArgs BUILD_PTX_FOR_ARCH) |
| 328 | set(multiValueArgs SRCS CUDA_ARCHS) |
| 329 | cmake_parse_arguments(arg "${options}" "${oneValueArgs}" |
| 330 | "${multiValueArgs}" ${ARGN} ) |
| 331 | |
| 332 | foreach(_ARCH ${arg_CUDA_ARCHS}) |
| 333 | # handle +PTX suffix: generate both sm and ptx codes if requested |
| 334 | string(FIND "${_ARCH}" "+PTX" _HAS_PTX) |
| 335 | if(NOT _HAS_PTX EQUAL -1) |
| 336 | string(REPLACE "+PTX" "" _BASE_ARCH "${_ARCH}") |
| 337 | string(REPLACE "." "" _STRIPPED_ARCH "${_BASE_ARCH}") |
| 338 | set_gencode_flag_for_srcs( |
| 339 | SRCS ${arg_SRCS} |
| 340 | ARCH "compute_${_STRIPPED_ARCH}" |
| 341 | CODE "sm_${_STRIPPED_ARCH}") |
| 342 | set_gencode_flag_for_srcs( |
| 343 | SRCS ${arg_SRCS} |
| 344 | ARCH "compute_${_STRIPPED_ARCH}" |
| 345 | CODE "compute_${_STRIPPED_ARCH}") |
| 346 | else() |
| 347 | string(REPLACE "." "" _STRIPPED_ARCH "${_ARCH}") |
| 348 | set_gencode_flag_for_srcs( |
| 349 | SRCS ${arg_SRCS} |
| 350 | ARCH "compute_${_STRIPPED_ARCH}" |
| 351 | CODE "sm_${_STRIPPED_ARCH}") |
| 352 | endif() |
| 353 | endforeach() |
| 354 | |
| 355 | if (${arg_BUILD_PTX_FOR_ARCH}) |
| 356 | list(SORT arg_CUDA_ARCHS COMPARE NATURAL ORDER ASCENDING) |
| 357 | list(GET arg_CUDA_ARCHS -1 _HIGHEST_ARCH) |
| 358 | if (_HIGHEST_ARCH VERSION_GREATER_EQUAL ${arg_BUILD_PTX_FOR_ARCH}) |
| 359 | string(REPLACE "." "" _PTX_ARCH "${arg_BUILD_PTX_FOR_ARCH}") |
| 360 | set_gencode_flag_for_srcs( |
| 361 | SRCS ${arg_SRCS} |
| 362 | ARCH "compute_${_PTX_ARCH}" |
| 363 | CODE "compute_${_PTX_ARCH}") |
| 364 | endif() |
| 365 | endif() |
| 366 | endmacro() |
| 367 | |
| 368 | # |
| 369 | # For the given `SRC_CUDA_ARCHS` list of gencode versions in the form |
| 370 | # `<major>.<minor>[letter]` compute the "loose intersection" with the |
| 371 | # `TGT_CUDA_ARCHS` list of gencodes. We also support the `+PTX` suffix in |
| 372 | # `SRC_CUDA_ARCHS` which indicates that the PTX code should be built when there |
| 373 | # is a CUDA_ARCH in `TGT_CUDA_ARCHS` that is equal to or larger than the |
| 374 | # architecture in `SRC_CUDA_ARCHS`. |
| 375 | # The loose intersection is defined as: |
| 376 | # { max{ x \in tgt | x <= y } | y \in src, { x \in tgt | x <= y } != {} } |
| 377 | # where `<=` is the version comparison operator. |
| 378 | # In other words, for each version in `TGT_CUDA_ARCHS` find the highest version |
| 379 | # in `SRC_CUDA_ARCHS` that is less or equal to the version in `TGT_CUDA_ARCHS`. |
| 380 | # We have special handling for x.0a, if x.0a is in `SRC_CUDA_ARCHS` and x.0 is |
| 381 | # in `TGT_CUDA_ARCHS` then we should remove x.0a from `SRC_CUDA_ARCHS` and add |
| 382 | # x.0a to the result (and remove x.0 from TGT_CUDA_ARCHS). |
| 383 | # The result is stored in `OUT_CUDA_ARCHS`. |
| 384 | # |
| 385 | # Example: |
| 386 | # SRC_CUDA_ARCHS="7.5;8.0;8.6;9.0;9.0a" |
| 387 | # TGT_CUDA_ARCHS="8.0;8.9;9.0" |
| 388 | # cuda_archs_loose_intersection(OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_ARCHS) |
| 389 | # OUT_CUDA_ARCHS="8.0;8.6;9.0;9.0a" |
| 390 | # |
| 391 | # Example With PTX: |
| 392 | # SRC_CUDA_ARCHS="8.0+PTX" |
| 393 | # TGT_CUDA_ARCHS="9.0" |
| 394 | # cuda_archs_loose_intersection(OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_ARCHS) |
| 395 | # OUT_CUDA_ARCHS="8.0+PTX" |
| 396 | # |
| 397 | function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_ARCHS) |
| 398 | set(_SRC_CUDA_ARCHS "${SRC_CUDA_ARCHS}") |
| 399 | set(_TGT_CUDA_ARCHS ${TGT_CUDA_ARCHS}) |
| 400 | |
| 401 | # handle +PTX suffix: separate base arch for matching, record PTX requests |
| 402 | set(_PTX_ARCHS) |
| 403 | foreach(_arch ${_SRC_CUDA_ARCHS}) |
| 404 | if(_arch MATCHES "\\+PTX$") |
| 405 | string(REPLACE "+PTX" "" _base "${_arch}") |
| 406 | list(APPEND _PTX_ARCHS "${_base}") |
| 407 | list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}") |
| 408 | list(APPEND _SRC_CUDA_ARCHS "${_base}") |
| 409 | endif() |
| 410 | endforeach() |
| 411 | list(REMOVE_DUPLICATES _PTX_ARCHS) |
| 412 | list(REMOVE_DUPLICATES _SRC_CUDA_ARCHS) |
| 413 | |
| 414 | # Handle architecture-specific suffixes (a/f) for SRC entries. |
| 415 | # First try exact base match (x.y), then cross-suffix match (x.ya / x.yf). |
| 416 | # For 'f' (family) suffix: if no exact/cross match, fall back to major-version |
| 417 | # match — e.g. SRC="12.0f" matches TGT="12.1a" since SM121 is in the SM12x |
| 418 | # family. The output uses TGT's value to preserve the user's compilation flags. |
| 419 | set(_CUDA_ARCHS) |
| 420 | # Resolve exact base matches before family fallbacks so a generic entry such |
| 421 | # as 10.0f cannot consume a 10.7 target that has a 10.7f source entry. |
| 422 | foreach(_arch ${_SRC_CUDA_ARCHS}) |
| 423 | if(_arch MATCHES "[af]$") |
| 424 | string(REGEX REPLACE "[af]$" "" _base "${_arch}") |
| 425 | if("${_base}" IN_LIST _TGT_CUDA_ARCHS) |
| 426 | list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}") |
| 427 | list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}") |
| 428 | list(APPEND _CUDA_ARCHS "${_arch}") |
| 429 | endif() |
| 430 | endif() |
| 431 | endforeach() |
| 432 | |
| 433 | foreach(_arch ${_SRC_CUDA_ARCHS}) |
| 434 | if(_arch MATCHES "[af]$") |
| 435 | list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}") |
| 436 | string(REGEX REPLACE "[af]$" "" _base "${_arch}") |
| 437 | if("${_base}a" IN_LIST _TGT_CUDA_ARCHS) |
| 438 | list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}a") |
| 439 | list(APPEND _CUDA_ARCHS "${_base}a") |
| 440 | elseif("${_base}f" IN_LIST _TGT_CUDA_ARCHS) |
| 441 | list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}f") |
| 442 | list(APPEND _CUDA_ARCHS "${_base}f") |
| 443 | elseif(_arch MATCHES "f$") |
| 444 | # Family suffix: match any TGT entry in the same major version family. |
| 445 | string(REGEX REPLACE "^([0-9]+)\\..*$" "\\1" _src_major "${_base}") |
| 446 | foreach(_tgt ${_TGT_CUDA_ARCHS}) |
| 447 | string(REGEX REPLACE "[af]$" "" _tgt_base "${_tgt}") |
| 448 | string(REGEX REPLACE "^([0-9]+)\\..*$" "\\1" _tgt_major "${_tgt_base}") |
| 449 | if(_tgt_major STREQUAL _src_major) |
| 450 | list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_tgt}") |
| 451 | list(APPEND _CUDA_ARCHS "${_tgt}") |
| 452 | break() |
| 453 | endif() |
| 454 | endforeach() |
| 455 | endif() |
| 456 | endif() |
| 457 | endforeach() |
| 458 | |
| 459 | # Symmetric handling: if TGT has x.ya/f and SRC has x.y (without suffix), |
| 460 | # preserve TGT's suffix in the output. |
| 461 | set(_tgt_copy ${_TGT_CUDA_ARCHS}) |
| 462 | foreach(_arch ${_tgt_copy}) |
| 463 | if(_arch MATCHES "[af]$") |
| 464 | string(REGEX REPLACE "[af]$" "" _base "${_arch}") |
| 465 | if ("${_base}" IN_LIST _SRC_CUDA_ARCHS) |
| 466 | list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_arch}") |
| 467 | list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_base}") |
| 468 | list(APPEND _CUDA_ARCHS "${_arch}") |
| 469 | endif() |
| 470 | endif() |
| 471 | endforeach() |
| 472 | |
| 473 | list(SORT _SRC_CUDA_ARCHS COMPARE NATURAL ORDER ASCENDING) |
| 474 | |
| 475 | # for each ARCH in TGT_CUDA_ARCHS find the highest arch in SRC_CUDA_ARCHS that |
| 476 | # is less or equal to ARCH (but has the same major version since SASS binary |
| 477 | # compatibility is only forward compatible within the same major version). |
| 478 | foreach(_ARCH ${_TGT_CUDA_ARCHS}) |
| 479 | set(_TMP_ARCH) |
| 480 | # Extract the major version of the target arch |
| 481 | string(REGEX REPLACE "^([0-9]+)\\..*$" "\\1" TGT_ARCH_MAJOR "${_ARCH}") |
| 482 | foreach(_SRC_ARCH ${_SRC_CUDA_ARCHS}) |
| 483 | # Extract the major version of the source arch |
| 484 | string(REGEX REPLACE "^([0-9]+)\\..*$" "\\1" SRC_ARCH_MAJOR "${_SRC_ARCH}") |
| 485 | # Check version-less-or-equal, and allow PTX arches to match across majors |
| 486 | if (_SRC_ARCH VERSION_LESS_EQUAL _ARCH) |
| 487 | if (_SRC_ARCH IN_LIST _PTX_ARCHS OR SRC_ARCH_MAJOR STREQUAL TGT_ARCH_MAJOR) |
| 488 | set(_TMP_ARCH "${_SRC_ARCH}") |
| 489 | endif() |
| 490 | else() |
| 491 | # If we hit a version greater than the target, we can break |
| 492 | break() |
| 493 | endif() |
| 494 | endforeach() |
| 495 | |
| 496 | # If we found a matching _TMP_ARCH, append it to _CUDA_ARCHS |
| 497 | if (_TMP_ARCH) |
| 498 | list(APPEND _CUDA_ARCHS "${_TMP_ARCH}") |
| 499 | endif() |
| 500 | endforeach() |
| 501 | |
| 502 | list(REMOVE_DUPLICATES _CUDA_ARCHS) |
| 503 | |
| 504 | # reapply +PTX suffix to architectures that requested PTX |
| 505 | set(_FINAL_ARCHS) |
| 506 | foreach(_arch ${_CUDA_ARCHS}) |
| 507 | if(_arch IN_LIST _PTX_ARCHS) |
| 508 | list(APPEND _FINAL_ARCHS "${_arch}+PTX") |
| 509 | else() |
| 510 | list(APPEND _FINAL_ARCHS "${_arch}") |
| 511 | endif() |
| 512 | endforeach() |
| 513 | set(_CUDA_ARCHS ${_FINAL_ARCHS}) |
| 514 | |
| 515 | set(${OUT_CUDA_ARCHS} ${_CUDA_ARCHS} PARENT_SCOPE) |
| 516 | endfunction() |
| 517 | |
| 518 | |
| 519 | function(cuda_archs_sm90plus OUT_CUDA_ARCHS TGT_CUDA_ARCHS) |
| 520 | if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) |
| 521 | cuda_archs_loose_intersection(_archs "9.0a;10.0f;10.7f;11.0f;12.0f" "${TGT_CUDA_ARCHS}") |
| 522 | else() |
| 523 | cuda_archs_loose_intersection(_archs "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${TGT_CUDA_ARCHS}") |
| 524 | endif() |
| 525 | set(${OUT_CUDA_ARCHS} ${_archs} PARENT_SCOPE) |
| 526 | endfunction() |
| 527 | |
| 528 | # |
| 529 | # Override the GPU architectures detected by cmake/torch and filter them by |
| 530 | # `GPU_SUPPORTED_ARCHES`. Sets the final set of architectures in |
| 531 | # `GPU_ARCHES`. This only applies to the HIP language since for CUDA we set |
| 532 | # the architectures on a per file basis. |
| 533 | # |
| 534 | # Note: this is defined as a macro since it updates `CMAKE_CUDA_FLAGS`. |
| 535 | # |
| 536 | macro(override_gpu_arches GPU_ARCHES GPU_LANG GPU_SUPPORTED_ARCHES) |
| 537 | set(_GPU_SUPPORTED_ARCHES_LIST ${GPU_SUPPORTED_ARCHES} ${ARGN}) |
| 538 | message(STATUS "${GPU_LANG} supported arches: ${_GPU_SUPPORTED_ARCHES_LIST}") |
| 539 | |
| 540 | if (${GPU_LANG} STREQUAL "HIP") |
| 541 | # |
| 542 | # `GPU_ARCHES` controls the `--offload-arch` flags. |
| 543 | # |
| 544 | # If PYTORCH_ROCM_ARCH env variable exists, then we take it as a list, |
| 545 | # if not, then we use CMAKE_HIP_ARCHITECTURES which was generated by calling |
| 546 | # "rocm_agent_enumerator" in "enable_language(HIP)" |
| 547 | # (in file Modules/CMakeDetermineHIPCompiler.cmake) |
| 548 | # |
| 549 | if(DEFINED ENV{PYTORCH_ROCM_ARCH}) |
| 550 | set(HIP_ARCHITECTURES $ENV{PYTORCH_ROCM_ARCH}) |
| 551 | else() |
| 552 | set(HIP_ARCHITECTURES ${CMAKE_HIP_ARCHITECTURES}) |
| 553 | endif() |
| 554 | # |
| 555 | # Find the intersection of the supported + detected architectures to |
| 556 | # set the module architecture flags. |
| 557 | # |
| 558 | set(${GPU_ARCHES}) |
| 559 | foreach (_ARCH ${HIP_ARCHITECTURES}) |
| 560 | if (_ARCH IN_LIST _GPU_SUPPORTED_ARCHES_LIST) |
| 561 | list(APPEND ${GPU_ARCHES} ${_ARCH}) |
| 562 | endif() |
| 563 | endforeach() |
| 564 | |
| 565 | if(NOT ${GPU_ARCHES}) |
| 566 | message(FATAL_ERROR |
| 567 | "None of the detected ROCm architectures: ${HIP_ARCHITECTURES} is" |
| 568 | " supported. Supported ROCm architectures are: ${_GPU_SUPPORTED_ARCHES_LIST}.") |
| 569 | endif() |
| 570 | endif() |
| 571 | endmacro() |
| 572 | |
| 573 | # |
| 574 | # Define a target named `MOD_NAME` for a single extension. The |
| 575 | # arguments are: |
| 576 | # |
| 577 | # DESTINATION <dest> - Module destination directory. |
| 578 | # LANGUAGE <lang> - The language for this module, e.g. CUDA, HIP, |
| 579 | # CXX, etc. |
| 580 | # SOURCES <sources> - List of source files relative to CMakeLists.txt |
| 581 | # directory. |
| 582 | # |
| 583 | # Optional arguments: |
| 584 | # |
| 585 | # ARCHITECTURES <arches> - A list of target architectures in cmake format. |
| 586 | # For GPU, refer to CMAKE_CUDA_ARCHITECTURES and |
| 587 | # CMAKE_HIP_ARCHITECTURES for more info. |
| 588 | # ARCHITECTURES will use cmake's defaults if |
| 589 | # not provided. |
| 590 | # COMPILE_FLAGS <flags> - Extra compiler flags passed to NVCC/hip. |
| 591 | # INCLUDE_DIRECTORIES <dirs> - Extra include directories. |
| 592 | # LIBRARIES <libraries> - Extra link libraries. |
| 593 | # WITH_SOABI - Generate library with python SOABI suffix name. |
| 594 | # USE_SABI <version> - Use python stable api <version> |
| 595 | # |
| 596 | # Note: optimization level/debug info is set via cmake build type. |
| 597 | # |
| 598 | function (define_extension_target MOD_NAME) |
| 599 | cmake_parse_arguments(PARSE_ARGV 1 |
| 600 | ARG |
| 601 | "WITH_SOABI" |
| 602 | "DESTINATION;LANGUAGE;USE_SABI" |
| 603 | "SOURCES;ARCHITECTURES;COMPILE_FLAGS;INCLUDE_DIRECTORIES;LIBRARIES") |
| 604 | |
| 605 | # Add hipify preprocessing step when building with HIP/ROCm. |
| 606 | if (ARG_LANGUAGE STREQUAL "HIP") |
| 607 | hipify_sources_target(ARG_SOURCES ${MOD_NAME} "${ARG_SOURCES}") |
| 608 | endif() |
| 609 | |
| 610 | if (ARG_WITH_SOABI) |
| 611 | set(SOABI_KEYWORD WITH_SOABI) |
| 612 | else() |
| 613 | set(SOABI_KEYWORD "") |
| 614 | endif() |
| 615 | |
| 616 | run_python(IS_FREETHREADED_PYTHON |
| 617 | "import sysconfig; print(1 if sysconfig.get_config_var(\"Py_GIL_DISABLED\") else 0)" |
| 618 | "Failed to determine whether interpreter is free-threaded") |
| 619 | |
| 620 | # Free-threaded Python doesn't yet support the stable ABI (see PEP 803/809), |
| 621 | # so avoid using the stable ABI under free-threading only. |
| 622 | if (ARG_USE_SABI AND NOT IS_FREETHREADED_PYTHON) |
| 623 | Python_add_library(${MOD_NAME} MODULE USE_SABI ${ARG_USE_SABI} ${SOABI_KEYWORD} "${ARG_SOURCES}") |
| 624 | else() |
| 625 | Python_add_library(${MOD_NAME} MODULE ${SOABI_KEYWORD} "${ARG_SOURCES}") |
| 626 | endif() |
| 627 | |
| 628 | if (ARG_LANGUAGE STREQUAL "HIP") |
| 629 | # Make this target dependent on the hipify preprocessor step. |
| 630 | add_dependencies(${MOD_NAME} hipify_all) |
| 631 | # Make sure we include the hipified versions of the headers, and avoid conflicts with the ones in the original source folder |
| 632 | target_include_directories(${MOD_NAME} PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/csrc |
| 633 | ${ARG_INCLUDE_DIRECTORIES}) |
| 634 | else() |
| 635 | target_include_directories(${MOD_NAME} PRIVATE csrc |
| 636 | ${ARG_INCLUDE_DIRECTORIES}) |
| 637 | endif() |
| 638 | |
| 639 | if (ARG_ARCHITECTURES) |
| 640 | set_target_properties(${MOD_NAME} PROPERTIES |
| 641 | ${ARG_LANGUAGE}_ARCHITECTURES "${ARG_ARCHITECTURES}") |
| 642 | endif() |
| 643 | |
| 644 | target_compile_options(${MOD_NAME} PRIVATE |
| 645 | $<$<COMPILE_LANGUAGE:${ARG_LANGUAGE}>:${ARG_COMPILE_FLAGS}>) |
| 646 | |
| 647 | target_compile_definitions(${MOD_NAME} PRIVATE |
| 648 | "-DTORCH_EXTENSION_NAME=${MOD_NAME}") |
| 649 | |
| 650 | target_link_libraries(${MOD_NAME} PRIVATE torch ${ARG_LIBRARIES}) |
| 651 | |
| 652 | # Don't use `TORCH_LIBRARIES` for CUDA since it pulls in a bunch of |
| 653 | # dependencies that are not necessary and may not be installed. |
| 654 | if (ARG_LANGUAGE STREQUAL "CUDA") |
| 655 | target_link_libraries(${MOD_NAME} PRIVATE torch CUDA::cudart CUDA::cuda_driver ${ARG_LIBRARIES}) |
| 656 | else() |
| 657 | target_link_libraries(${MOD_NAME} PRIVATE torch ${TORCH_LIBRARIES} ${ARG_LIBRARIES}) |
| 658 | endif() |
| 659 | |
| 660 | install(TARGETS ${MOD_NAME} LIBRARY DESTINATION ${ARG_DESTINATION} COMPONENT ${MOD_NAME}) |
| 661 | endfunction() |