cmake/TorchConfig.cmake.in

# FindTorch
# -------
#
# Finds the Torch library
#
# This will define the following variables:
#
#   TORCH_FOUND        -- True if the system has the Torch library
#   TORCH_INCLUDE_DIRS -- The include directories for torch
#   TORCH_LIBRARIES    -- Libraries to link to
#
# and the following imported targets:
#
#   Torch
#
# and the following functions:
#
#   torch_add_custom_op_library(<name> <source_files>)

if ($ENV{TORCH_INSTALL_PREFIX})
  set(TORCH_INSTALL_PREFIX $ENV{TORCH_INSTALL_PREFIX})
else()
  # Assume we are in <install-prefix>/share/cmake/Torch/TorchConfig.cmake
  get_filename_component(CMAKE_CURRENT_LIST_DIR "${CMAKE_CURRENT_LIST_FILE}" PATH)
  get_filename_component(TORCH_INSTALL_PREFIX "${CMAKE_CURRENT_LIST_DIR}/../../../" ABSOLUTE)
endif()

# Include directories.
set(TORCH_INCLUDE_DIRS "${TORCH_INSTALL_PREFIX}/lib/include")

# Library dependencies.
find_package(Caffe2 REQUIRED)
find_library(TORCH_LIBRARY torch PATHS "${TORCH_INSTALL_PREFIX}/lib")

set(TORCH_LIBRARIES ${TORCH_LIBRARY} ${Caffe2_MAIN_LIBS})
if (@USE_CUDA@)
  if(MSVC)
    set(NVTOOLEXT_HOME "C:/Program Files/NVIDIA Corporation/NvToolsExt")
    if ($ENV{NVTOOLEXT_HOME})
      set(NVTOOLEXT_HOME $ENV{NVTOOLEXT_HOME})
    endif()
    set(TORCH_CUDA_LIBRARIES
      ${NVTOOLEXT_HOME}/lib/x64/nvToolsExt64_1.lib
      ${CUDA_LIBRARIES})
    list(APPEND TORCH_INCLUDE_DIRS "${NVTOOLEXT_HOME}/include")
  elseif(APPLE)
    set(TORCH_CUDA_LIBRARIES
      ${CUDA_TOOLKIT_ROOT_DIR}/lib/libcudart.dylib
      ${CUDA_TOOLKIT_ROOT_DIR}/lib/libnvrtc.dylib
      ${CUDA_TOOLKIT_ROOT_DIR}/lib/libnvToolsExt.dylib
      ${CUDA_LIBRARIES})
  else()
    set(TORCH_CUDA_LIBRARIES
      ${CUDA_CUDA_LIB}
      ${CUDA_NVRTC_LIB}
      ${CUDA_TOOLKIT_ROOT_DIR}/lib64/libnvToolsExt.so
      ${CUDA_LIBRARIES})
  endif()
  list(APPEND TORCH_LIBRARIES ${TORCH_CUDA_LIBRARIES})
endif()

# Creates a shared library <name> with the correct include directories
# and linker flags set to include Torch header files and link with Torch
# libraries. Also sets the C++ standard version to C++11. All options
# can be override by specifying further options on the `<name>` CMake target.
function(torch_add_custom_op_library name source_files)
  add_library(${name} SHARED ${source_files})
  target_include_directories(${name} PUBLIC "${TORCH_INCLUDE_DIRS}")
  target_link_libraries(${name} "${TORCH_LIBRARIES}")
  set_property(TARGET ${name} PROPERTY CXX_STANDARD 11)
endfunction(torch_add_custom_op_library)