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/Acts
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 endif()
0019
0020 if(ACTS_GNN_ENABLE_MODULEMAP)
0021     target_compile_definitions(ActsPluginGnn PUBLIC ACTS_GNN_WITH_MODULEMAP)
0022
0023     target_link_libraries(
0024         ActsPluginGnn
0025         PRIVATE ModuleMapGraph::CPU ModuleMapGraph::GPU
0026     )
0027     target_sources(ActsPluginGnn PRIVATE src/ModuleMapCuda.cu)
0028 endif()
0029
0030 if(ACTS_GNN_ENABLE_ONNX)
0031     target_sources(ActsPluginGnn PRIVATE src/OnnxEdgeClassifier.cpp)
0032     target_link_libraries(ActsPluginGnn PRIVATE onnxruntime::onnxruntime)
0033     target_compile_definitions(ActsPluginGnn PUBLIC ACTS_GNN_ONNX_BACKEND)
0034 endif()
0035
0036 if(ACTS_GNN_ENABLE_TORCH)
0037     target_sources(
0038         ActsPluginGnn
0039         PRIVATE
0040             src/TorchEdgeClassifier.cpp
0041             src/TorchMetricLearning.cpp
0042             src/buildEdges.cpp
0043     )
0044 endif()
0045
0046 if(ACTS_GNN_ENABLE_TENSORRT)
0047     find_package(TensorRT REQUIRED)
0048     message(STATUS "Found TensorRT ${TensorRT_VERSION}")
0049     target_link_libraries(ActsPluginGnn PUBLIC trt::nvinfer trt::nvinfer_plugin)
0050     target_sources(ActsPluginGnn PRIVATE src/TensorRTEdgeClassifier.cpp)
0051     target_compile_definitions(ActsPluginGnn PUBLIC ACTS_GNN_WITH_TENSORRT)
0052 endif()
0053
0054 target_include_directories(
0055     ActsPluginGnn
0056     PUBLIC
0057         $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>
0058         $<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>
0059 )
0060
0061 target_link_libraries(ActsPluginGnn PUBLIC Acts::Core Boost::boost)
0062
0063 target_link_libraries(ActsPluginGnn PRIVATE std::filesystem)
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} frnn)
0092     find_package(TorchScatter QUIET)
0093     if(NOT TARGET TorchScatter::TorchScatter)
0094         message(
0095             WARNING
0096             "Torch scatter not found, models that rely on torch-scatter will not work"
0097         )
0098     else()
0099         target_link_libraries(ActsPluginGnn PRIVATE TorchScatter::TorchScatter)
0100         # Should not discard TorchScatter even if its not needed at this point
0101         # since we need the scatter_max operation in the torch script later
0102         target_link_options(ActsPluginGnn PUBLIC "-Wl,-no-as-needed")
0103     endif()
0104 endif()