CMakeLists.txt

cmake_minimum_required(VERSION 3.14)
project(autoware_tensorrt_plugins)

find_package(autoware_cmake REQUIRED)
autoware_package()

add_compile_options(-Wno-deprecated-declarations)

option(CUDA_VERBOSE "Verbose output of CUDA modules" OFF)

# set flags for CUDA availability
option(CUDA_AVAIL "CUDA available" OFF)
find_package(CUDA)
if(CUDA_FOUND)
  find_library(CUBLAS_LIBRARIES cublas HINTS
    ${CUDA_TOOLKIT_ROOT_DIR}/lib64
    ${CUDA_TOOLKIT_ROOT_DIR}/lib
  )
  if(CUDA_VERBOSE)
    message("CUDA is available!")
    message("CUDA Libs: ${CUDA_LIBRARIES}")
    message("CUDA Headers: ${CUDA_INCLUDE_DIRS}")
  endif()
  # Note: cublas_device was depreciated in CUDA version 9.2
  #       https://forums.developer.nvidia.com/t/where-can-i-find-libcublas-device-so-or-libcublas-device-a/67251/4
  #       In LibTorch, CUDA_cublas_device_LIBRARY is used.
  unset(CUDA_cublas_device_LIBRARY CACHE)
  set(CUDA_AVAIL ON)
else()
  message("CUDA NOT FOUND")
  set(CUDA_AVAIL OFF)
endif()

# set flags for TensorRT availability
option(TRT_AVAIL "TensorRT available" OFF)
# try to find the tensorRT modules
find_library(NVINFER nvinfer)
find_library(NVONNXPARSER nvonnxparser)
if(NVINFER AND NVONNXPARSER)
  if(CUDA_VERBOSE)
    message("TensorRT is available!")
    message("NVINFER: ${NVINFER}")
    message("NVONNXPARSER: ${NVONNXPARSER}")
  endif()
  set(TRT_AVAIL ON)
else()
  message("TensorRT is NOT available")
  set(TRT_AVAIL OFF)
endif()

# set flags for spconv availability
option(SPCONV_AVAIL "spconv available" OFF)
# try to find spconv
find_package(cumm)
find_package(spconv)
if(${cumm_FOUND} AND ${spconv_FOUND})
  message(STATUS "spconv is available!")
  set(SPCONV_AVAIL ON)
else()
  message(WARNING "spconv is NOT available")
  set(SPCONV_AVAIL OFF)
endif()

if(TRT_AVAIL AND CUDA_AVAIL AND SPCONV_AVAIL)
  find_package(tensorrt_cmake_module REQUIRED)
  find_package(TENSORRT)
  if(TENSORRT_VERSION VERSION_LESS 10.0)
    message(WARNING "Unsupported version TensorRT ${TENSORRT_VERSION} detected. This package requires TensorRT 10.0 or later.")
    return()
  endif()

  find_package(ament_cmake_auto REQUIRED)
  ament_auto_find_build_dependencies()

  include_directories(
    include
    ${CUDA_INCLUDE_DIRS}
  )

  if(CMAKE_BUILD_TYPE STREQUAL "Debug")
    set(CMAKE_CUDA_FLAGS ${CMAKE_CUDA_FLAGS} "-g -G")
  endif()

  add_definitions("-DTV_CUDA")

  # cSpell:ignore expt
  list(APPEND CUDA_NVCC_FLAGS "--expt-relaxed-constexpr -diag-suppress 1675 --extended-lambda")
  list(APPEND CUDA_NVCC_FLAGS "-gencode arch=compute_75,code=sm_75")
  list(APPEND CUDA_NVCC_FLAGS "-gencode arch=compute_86,code=sm_86")
  list(APPEND CUDA_NVCC_FLAGS "-gencode arch=compute_87,code=sm_87")
  list(APPEND CUDA_NVCC_FLAGS "-gencode arch=compute_89,code=sm_89")
  if(CUDA_VERSION VERSION_GREATER_EQUAL "12.8")
    if(CUDA_VERSION VERSION_LESS "13.0")
      list(APPEND CUDA_NVCC_FLAGS "-gencode arch=compute_101,code=sm_101")
    else()  # CUDA 13.0 renamed SM101 to SM110
      list(APPEND CUDA_NVCC_FLAGS "-gencode arch=compute_110,code=sm_110")
    endif()
    list(APPEND CUDA_NVCC_FLAGS "-gencode arch=compute_120,code=sm_120")
    list(APPEND CUDA_NVCC_FLAGS "-gencode arch=compute_120,code=compute_120")
  else()
    list(APPEND CUDA_NVCC_FLAGS "-gencode arch=compute_89,code=compute_89")
  endif()

  cuda_add_library(cuda_ops SHARED
    src/argsort_ops/argsort.cu
    src/bev_ops/bev_pool_cuda.cu
    src/scatter_ops/segment_csr.cu
    src/unique_ops/unique.cu
    src/multi_scale_deform_attn_ops/ms_deform_attn_kernel.cu
    src/rotate_ops/rotate_kernel.cu
    src/select_and_pad_ops/select_and_pad_kernel.cu
  )

  add_library(${PROJECT_NAME} SHARED
    src/argsort_plugin.cpp
    src/argsort_plugin_creator.cpp
    src/quick_cumsum_cuda_plugin.cpp
    src/quick_cumsum_cuda_plugin_creator.cpp
    src/get_indices_pairs_implicit_gemm_plugin_creator.cpp
    src/get_indices_pairs_implicit_gemm_plugin.cpp
    src/get_indices_pairs_plugin_creator.cpp
    src/get_indices_pairs_plugin.cpp
    src/implicit_gemm_plugin_creator.cpp
    src/implicit_gemm_plugin.cpp
    src/indice_conv_plugin_creator.cpp # cSpell:ignore indice
    src/indice_conv_plugin.cpp
    src/multi_scale_deformable_attention_plugin.cpp
    src/multi_scale_deformable_attention_plugin_creator.cpp
    src/rotate_plugin.cpp
    src/rotate_plugin_creator.cpp
    src/segment_csr_plugin_creator.cpp
    src/segment_csr_plugin.cpp
    src/select_and_pad_plugin.cpp
    src/select_and_pad_plugin_creator.cpp
    src/unique_plugin.cpp
    src/unique_plugin_creator.cpp
    src/plugin_registration.cpp
    src/plugin_utils.cpp
  )

  target_compile_definitions(${PROJECT_NAME} PRIVATE _GLIBCXX_USE_CXX11_ABI=1)

  target_link_libraries(${PROJECT_NAME} PRIVATE
      ${NVINFER}
      CUDA::cudart
      cuda_ops
      spconv::spconv
  )

  if(BUILD_TESTING)
    find_package(ament_cmake_gtest REQUIRED)

    ament_add_gtest(reference_kernels_test
      test/argsort_kernel_test.cpp
      test/unique_kernel_test.cpp
    )
    if(TARGET reference_kernels_test)
      target_link_libraries(reference_kernels_test
        CUDA::cudart
        cuda_ops
      )
      target_include_directories(reference_kernels_test PRIVATE
        include
        test
        ${autoware_cuda_utils_INCLUDE_DIRS}
        ${CUDA_INCLUDE_DIRS}
      )
      target_compile_definitions(reference_kernels_test PRIVATE _GLIBCXX_USE_CXX11_ABI=1)
    endif()
  endif()

  install(
    TARGETS ${PROJECT_NAME}
    DESTINATION share/${PROJECT_NAME}/plugins
  )

  install(
    TARGETS cuda_ops
    LIBRARY DESTINATION lib
  )

  ament_auto_package()
else()
  find_package(ament_cmake_auto REQUIRED)
  ament_auto_find_build_dependencies()
endif()