Plugins/Gnn/CMakeLists.txt

0001 acts_add_library(
0002     PluginGnn
0003     src/GnnPipeline.cpp
0004     src/Tensor.cpp
0005     src/BoostTrackBuilding.cpp
0006     src/TruthGraphMetricsHook.cpp
0007     src/GraphStoreHook.cpp
0008     ACTS_INCLUDE_FOLDER include/ActsPlugins
0009 )
0010
0011 if(ACTS_GNN_ENABLE_CUDA)
0012     target_compile_definitions(ActsPluginGnn PUBLIC ACTS_GNN_WITH_CUDA)
0013     target_sources(
0014         ActsPluginGnn
0015         PRIVATE src/CudaTrackBuilding.cu src/JunctionRemoval.cu src/Tensor.cu
0016     )
0017     target_link_libraries(ActsPluginGnn PUBLIC CUDA::nvtx3)
0018     target_compile_definitions(ActsPluginGnn PUBLIC ACTS_GNN_NVTX)
0019 endif()
0020
0021 if(ACTS_GNN_ENABLE_MODULEMAP)
0022     target_compile_definitions(ActsPluginGnn PUBLIC ACTS_GNN_WITH_MODULEMAP)
0023
0024     target_link_libraries(
0025         ActsPluginGnn
0026         PRIVATE ModuleMapGraph::CPU ModuleMapGraph::GPU
0027     )
0028     target_sources(ActsPluginGnn PRIVATE src/ModuleMapCuda.cu)
0029 endif()
0030
0031 if(ACTS_GNN_ENABLE_ONNX)
0032     target_sources(ActsPluginGnn PRIVATE src/OnnxEdgeClassifier.cpp)
0033     target_link_libraries(ActsPluginGnn PRIVATE onnxruntime::onnxruntime)
0034     target_compile_definitions(ActsPluginGnn PUBLIC ACTS_GNN_ONNX_BACKEND)
0035 endif()
0036
0037 if(ACTS_GNN_ENABLE_TORCH)
0038     target_sources(
0039         ActsPluginGnn
0040         PRIVATE
0041             src/TorchEdgeClassifier.cpp
0042             src/TorchMetricLearning.cpp
0043             src/buildEdges.cpp
0044     )
0045 endif()
0046
0047 if(ACTS_GNN_ENABLE_TENSORRT)
0048     find_package(TensorRT REQUIRED)
0049     message(STATUS "Found TensorRT ${TensorRT_VERSION}")
0050     target_link_libraries(ActsPluginGnn PUBLIC trt::nvinfer trt::nvinfer_plugin)
0051     target_sources(ActsPluginGnn PRIVATE src/TensorRTEdgeClassifier.cpp)
0052     target_compile_definitions(ActsPluginGnn PUBLIC ACTS_GNN_WITH_TENSORRT)
0053 endif()
0054
0055 target_include_directories(
0056     ActsPluginGnn
0057     PUBLIC
0058         $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>
0059         $<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>
0060 )
0061
0062 target_link_libraries(ActsPluginGnn PUBLIC Acts::Core Boost::boost)
0063
0064 if(ACTS_GNN_ENABLE_CUDA)
0065     target_compile_features(ActsPluginGnn PUBLIC cuda_std_20)
0066     set_target_properties(
0067         ActsPluginGnn
0068         PROPERTIES CUDA_STANDARD_REQUIRED ON CUDA_SEPARABLE_COMPILATION OFF
0069     )
0070     target_compile_options(
0071         ActsPluginGnn
0072         PRIVATE
0073             $<$<COMPILE_LANGUAGE:CUDA>:
0074             --dopt=on
0075             --generate-line-info
0076             --expt-relaxed-constexpr
0077             --extended-lambda>
0078     )
0079     target_compile_definitions(
0080         ActsPluginGnn
0081         PUBLIC CUDA_API_PER_THREAD_DEFAULT_STREAM
0082     )
0083     target_link_libraries(ActsPluginGnn PUBLIC CUDA::cudart)
0084 else()
0085     target_compile_definitions(ActsPluginGnn PUBLIC ACTS_GNN_CPUONLY)
0086 endif()
0087
0088 if(ACTS_GNN_ENABLE_TORCH)
0089     target_compile_definitions(ActsPluginGnn PUBLIC ACTS_GNN_TORCH_BACKEND)
0090     target_link_libraries(ActsPluginGnn PRIVATE ${TORCH_LIBRARIES})
0091     if(ACTS_GNN_ENABLE_CUDA)
0092         target_link_libraries(ActsPluginGnn PRIVATE frnn)
0093     endif()
0094     find_package(TorchScatter QUIET)
0095     if(NOT TARGET TorchScatter::TorchScatter)
0096         message(
0097             WARNING
0098             "Torch scatter not found, models that rely on torch-scatter will not work"
0099         )
0100     else()
0101         target_link_libraries(ActsPluginGnn PRIVATE TorchScatter::TorchScatter)
0102         # Should not discard TorchScatter even if its not needed at this point
0103         # since we need the scatter_max operation in the torch script later
0104         target_link_options(ActsPluginGnn PUBLIC "-Wl,-no-as-needed")
0105     endif()
0106 endif()