Plugins/Gnn/CMakeLists.txt

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