Files
nvidia--tensorrt/cmake/modules/ShouldCompileKernel.cmake
T
wehub-resource-sync c8a779b1bb
Docker Image CI / build-ubuntu2004 (push) Waiting to run
chore: import upstream snapshot with attribution
2026-07-13 13:36:55 +08:00

106 lines
4.4 KiB
CMake

# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# \brief Converts a SM string (i.e. 86+abc) into the numeric SM version (i.e. 86).
# \returns the sm in the name specified by OUT_VAR.
function(get_numeric_sm SM OUT_VAR)
# Convert the SM string to a numeric value
if(${SM} MATCHES "^([0-9]+).*$")
set(${OUT_VAR} ${CMAKE_MATCH_1} PARENT_SCOPE)
else()
message(FATAL_ERROR "Invalid SM version: ${SM}")
endif()
endfunction()
# \brief Converts the CMAKE_CUDA_ARCHITECTURES list into a list of numeric SM values.
# \returns the list in the name specified by OUT_VAR.
function(get_all_numeric_sms OUT_VAR)
set(ALL_NUMERIC_SMS "")
foreach(SM IN LISTS CMAKE_CUDA_ARCHITECTURES)
get_numeric_sm(${SM} "SM")
list(APPEND ALL_NUMERIC_SMS ${SM})
endforeach()
set(${OUT_VAR} ${ALL_NUMERIC_SMS} PARENT_SCOPE)
endfunction()
# \brief Converts the list returned by get_all_numeric_sms into a list of arch values.
# \returns the list in the name specified by OUT_VAR for native platform and OUT_VAR_CROSS for cross OS support. e.g. ptx, sm75, sm80, sm86, sm89, sm100, sm120.
function(get_all_fatbin_archs OUT_VAR OUT_VAR_CROSS)
# Use get_all_numeric_sms to get SM values and convert them to sm-prefixed format
set(ARCH_LIST "")
set(ARCH_LIST_CROSS "")
get_all_numeric_sms(NUMERIC_SMS)
foreach(SM IN LISTS NUMERIC_SMS)
list(APPEND ARCH_LIST "sm${SM}")
endforeach()
# Note: sm89 it is missing in NUMERIC_SMS since TRT treats sm89 as sm86.
# We should add sm89 to the list to generate the builder resource for sm89.
# If only sm86 is in the list, it means this build only supports sm86,
# so no need to add sm89.
list(FIND ARCH_LIST "sm86" SM86_INDEX)
list(FIND ARCH_LIST "sm89" SM89_INDEX)
list(LENGTH ARCH_LIST ARCH_LIST_COUNT)
if(${SM86_INDEX} GREATER_EQUAL 0 AND ${SM89_INDEX} EQUAL -1 AND ${ARCH_LIST_COUNT} GREATER 1)
list(APPEND ARCH_LIST "sm89")
endif()
# On SBSA exclude sm87 explicitly since it is proxied to sm86.
if(${TRT_IS_ARM_SERVER})
list(FILTER ARCH_LIST EXCLUDE REGEX "sm87")
endif()
# Exclude sm103 explicitly since it is proxied to sm100.
list(FILTER ARCH_LIST EXCLUDE REGEX "sm103")
# Exclude sm121 explicitly since it is proxied to sm120.
list(FILTER ARCH_LIST EXCLUDE REGEX "sm121")
# There is also a klib which only contains PTX code.
list(APPEND ARCH_LIST "ptx")
set(ARCH_LIST_CROSS ${ARCH_LIST})
set(${OUT_VAR} ${ARCH_LIST} PARENT_SCOPE)
set(${OUT_VAR_CROSS} ${ARCH_LIST_CROSS} PARENT_SCOPE)
endfunction()
# Certain cubins are binary compatible between different SM versions, so they are reused.
# This function checks if a SM-named file should be compiled based on current SM enablement.
# Specifically, the SM80 files are compiled if either 80, 86, or 89 are enabled.
function(should_compile_kernel SM OUT_VAR)
get_all_numeric_sms(__TRT_NUMERIC_CUDA_ARCHS)
# If the target SM is any of 80/86/89, we need to check if any of those are enabled in __TRT_NUMERIC_CUDA_ARCHS.
if((${SM} EQUAL 80) OR (${SM} EQUAL 86) OR (${SM} EQUAL 89))
list(FIND __TRT_NUMERIC_CUDA_ARCHS 80 SM80_INDEX)
list(FIND __TRT_NUMERIC_CUDA_ARCHS 86 SM86_INDEX)
list(FIND __TRT_NUMERIC_CUDA_ARCHS 89 SM89_INDEX)
if((NOT ${SM80_INDEX} EQUAL -1) OR
(NOT ${SM86_INDEX} EQUAL -1) OR
(NOT ${SM89_INDEX} EQUAL -1)
)
set(${OUT_VAR} TRUE PARENT_SCOPE)
else()
set(${OUT_VAR} FALSE PARENT_SCOPE)
endif()
else()
list(FIND __TRT_NUMERIC_CUDA_ARCHS ${SM} SM_INDEX)
if (NOT ${SM_INDEX} EQUAL -1)
set(${OUT_VAR} TRUE PARENT_SCOPE)
else()
set(${OUT_VAR} FALSE PARENT_SCOPE)
endif()
endif()
endfunction()