Add AMD support via ort alternative to trt
This commit is contained in:
+162
-55
@@ -21,46 +21,162 @@ add_subdirectory(external/KPN)
|
||||
find_package(OpenCV 4 REQUIRED COMPONENTS
|
||||
core imgproc imgcodecs videoio dnn objdetect highgui)
|
||||
|
||||
# ONNX Runtime (SCRFD detector — cv::dnn cannot handle dynamic Shape nodes)
|
||||
find_library(ORT_LIB onnxruntime REQUIRED
|
||||
HINTS /usr/lib /usr/local/lib)
|
||||
find_path(ORT_INCLUDE onnxruntime_cxx_api.h
|
||||
PATH_SUFFIXES onnxruntime
|
||||
HINTS /usr/include /usr/local/include
|
||||
REQUIRED)
|
||||
add_library(onnxruntime UNKNOWN IMPORTED)
|
||||
set_target_properties(onnxruntime PROPERTIES
|
||||
IMPORTED_LOCATION "${ORT_LIB}"
|
||||
INTERFACE_INCLUDE_DIRECTORIES "${ORT_INCLUDE}")
|
||||
message(STATUS "ONNX Runtime: ${ORT_LIB} headers: ${ORT_INCLUDE}")
|
||||
# ── Model paths ───────────────────────────────────────────────────────────────
|
||||
# Defined early so the backend object libraries below can embed it.
|
||||
set(SAE_MODELS_DIR "${CMAKE_SOURCE_DIR}/models"
|
||||
CACHE PATH "Directory containing ONNX model files")
|
||||
|
||||
# TensorRT + CUDA runtime (raw-TRT ArcFace embedder; activated by --arcface-engine).
|
||||
find_library(NVINFER_LIB nvinfer
|
||||
HINTS /usr/lib /usr/local/lib /opt/tensorrt/lib)
|
||||
find_path(NVINFER_INCLUDE NvInfer.h
|
||||
HINTS /usr/include /usr/local/include /opt/tensorrt/include)
|
||||
find_library(CUDART_LIB cudart
|
||||
HINTS /opt/cuda/lib64 /usr/local/cuda/lib64 /usr/lib)
|
||||
find_path(CUDART_INCLUDE cuda_runtime_api.h
|
||||
HINTS /opt/cuda/targets/x86_64-linux/include /opt/cuda/include
|
||||
/usr/local/cuda/include /usr/include)
|
||||
find_library(CUBLAS_LIB cublas
|
||||
HINTS /opt/cuda/targets/x86_64-linux/lib /opt/cuda/lib64
|
||||
/usr/local/cuda/lib64 /usr/lib)
|
||||
if(NOT (NVINFER_LIB AND NVINFER_INCLUDE AND CUDART_LIB AND CUDART_INCLUDE AND CUBLAS_LIB))
|
||||
message(FATAL_ERROR
|
||||
"TensorRT or CUDA runtime not found "
|
||||
"(nvinfer=${NVINFER_LIB} headers=${NVINFER_INCLUDE} "
|
||||
"cudart=${CUDART_LIB} cublas=${CUBLAS_LIB} headers=${CUDART_INCLUDE})")
|
||||
# ── Backend selection ─────────────────────────────────────────────────────────
|
||||
# Two independent compile-time axes. The core application is agnostic to both:
|
||||
# only the matching backend .cpp (in src/backends/) is compiled, and the backend
|
||||
# headers (onnxruntime / NvInfer.h / cublas / rocblas) never reach core TUs.
|
||||
#
|
||||
# SAE_INFERENCE_BACKEND ORT → SCRFD + ArcFace via ONNX Runtime (.onnx models)
|
||||
# TRT → SCRFD + ArcFace via raw TensorRT (.engine files)
|
||||
# SAE_GEMM_BACKEND ROCM → gallery similarity GEMM via rocBLAS / HIP
|
||||
# CUDA → gallery similarity GEMM via cuBLAS / CUDA
|
||||
set(SAE_INFERENCE_BACKEND "ORT" CACHE STRING "Inference backend: ORT | TRT")
|
||||
set(SAE_GEMM_BACKEND "ROCM" CACHE STRING "Gallery GEMM backend: ROCM | CUDA")
|
||||
set_property(CACHE SAE_INFERENCE_BACKEND PROPERTY STRINGS ORT TRT)
|
||||
set_property(CACHE SAE_GEMM_BACKEND PROPERTY STRINGS ROCM CUDA)
|
||||
|
||||
# Enable the ORT TensorRT/CUDA execution providers inside the ORT inference
|
||||
# backend (only meaningful when ORT was built with the TensorRT EP). Off by
|
||||
# default so ROCm/CPU builds don't reference unavailable EPs.
|
||||
option(SAE_ORT_TRT_EP "ORT backend: enable TensorRT/CUDA execution providers" OFF)
|
||||
|
||||
# Back-compat: a legacy -DSAE_WITH_TRT=ON/OFF seeds the new vars (ON⇒TRT+CUDA,
|
||||
# OFF⇒ORT+ROCM) unless the user set them explicitly.
|
||||
if(DEFINED SAE_WITH_TRT)
|
||||
if(SAE_WITH_TRT)
|
||||
set(SAE_INFERENCE_BACKEND "TRT" CACHE STRING "" FORCE)
|
||||
set(SAE_GEMM_BACKEND "CUDA" CACHE STRING "" FORCE)
|
||||
else()
|
||||
set(SAE_INFERENCE_BACKEND "ORT" CACHE STRING "" FORCE)
|
||||
set(SAE_GEMM_BACKEND "ROCM" CACHE STRING "" FORCE)
|
||||
endif()
|
||||
message(STATUS "SAE_WITH_TRT=${SAE_WITH_TRT} (legacy) → "
|
||||
"SAE_INFERENCE_BACKEND=${SAE_INFERENCE_BACKEND} "
|
||||
"SAE_GEMM_BACKEND=${SAE_GEMM_BACKEND}")
|
||||
endif()
|
||||
add_library(trt_runtime INTERFACE)
|
||||
target_include_directories(trt_runtime INTERFACE
|
||||
"${NVINFER_INCLUDE}" "${CUDART_INCLUDE}")
|
||||
target_link_libraries(trt_runtime INTERFACE
|
||||
"${NVINFER_LIB}" "${CUDART_LIB}" "${CUBLAS_LIB}")
|
||||
message(STATUS "TensorRT: ${NVINFER_LIB} CUDA runtime: ${CUDART_LIB} cuBLAS: ${CUBLAS_LIB}")
|
||||
|
||||
# FFmpeg (NVDEC hardware video decode + swscale colour conversion)
|
||||
if(NOT SAE_INFERENCE_BACKEND MATCHES "^(ORT|TRT)$")
|
||||
message(FATAL_ERROR "SAE_INFERENCE_BACKEND must be ORT or TRT (got '${SAE_INFERENCE_BACKEND}')")
|
||||
endif()
|
||||
if(NOT SAE_GEMM_BACKEND MATCHES "^(ROCM|CUDA)$")
|
||||
message(FATAL_ERROR "SAE_GEMM_BACKEND must be ROCM or CUDA (got '${SAE_GEMM_BACKEND}')")
|
||||
endif()
|
||||
|
||||
# CUDA runtime is needed by both TRT inference and CUDA GEMM — find it once.
|
||||
function(sae_find_cudart)
|
||||
if(TARGET cudart_dep)
|
||||
return()
|
||||
endif()
|
||||
find_library(CUDART_LIB cudart
|
||||
HINTS /opt/cuda/lib64 /usr/local/cuda/lib64 /usr/lib)
|
||||
find_path(CUDART_INCLUDE cuda_runtime_api.h
|
||||
HINTS /opt/cuda/targets/x86_64-linux/include /opt/cuda/include
|
||||
/usr/local/cuda/include /usr/include)
|
||||
if(NOT (CUDART_LIB AND CUDART_INCLUDE))
|
||||
message(FATAL_ERROR "CUDA runtime not found (cudart=${CUDART_LIB} headers=${CUDART_INCLUDE}).")
|
||||
endif()
|
||||
add_library(cudart_dep INTERFACE)
|
||||
target_include_directories(cudart_dep INTERFACE "${CUDART_INCLUDE}")
|
||||
target_link_libraries(cudart_dep INTERFACE "${CUDART_LIB}")
|
||||
set_property(GLOBAL PROPERTY sae_cudart_found TRUE)
|
||||
endfunction()
|
||||
|
||||
# ── Inference backend dependency: builds the `inference_backend` object lib ────
|
||||
if(SAE_INFERENCE_BACKEND STREQUAL "ORT")
|
||||
find_library(ORT_LIB onnxruntime REQUIRED
|
||||
HINTS /usr/lib64/rocm/lib /usr/lib /usr/local/lib)
|
||||
find_path(ORT_INCLUDE onnxruntime_cxx_api.h
|
||||
PATH_SUFFIXES onnxruntime
|
||||
HINTS /usr/lib64/rocm/include/onnxruntime /usr/include/onnxruntime /usr/local/include/onnxruntime
|
||||
/usr/lib64/rocm/include /usr/include /usr/local/include
|
||||
REQUIRED)
|
||||
# The include directive is <onnxruntime/onnxruntime_cxx_api.h>, so we need the
|
||||
# parent of the onnxruntime/ subdirectory on the include path.
|
||||
get_filename_component(ORT_INCLUDE_PARENT "${ORT_INCLUDE}" DIRECTORY)
|
||||
if(NOT EXISTS "${ORT_INCLUDE_PARENT}/onnxruntime")
|
||||
set(ORT_INCLUDE_PARENT "${ORT_INCLUDE}")
|
||||
endif()
|
||||
message(STATUS "Inference backend: ORT (${ORT_LIB} headers: ${ORT_INCLUDE_PARENT})")
|
||||
|
||||
add_library(inference_backend OBJECT src/backends/ort_backend.cpp)
|
||||
set_target_properties(inference_backend PROPERTIES POSITION_INDEPENDENT_CODE ON)
|
||||
target_include_directories(inference_backend PRIVATE src "${ORT_INCLUDE_PARENT}")
|
||||
target_link_libraries(inference_backend PRIVATE ${OpenCV_LIBS} "${ORT_LIB}")
|
||||
target_compile_definitions(inference_backend PRIVATE
|
||||
SAE_MODELS_DIR="${SAE_MODELS_DIR}"
|
||||
$<$<BOOL:${SAE_ORT_TRT_EP}>:SAE_ORT_WITH_TRT_EP>)
|
||||
else() # TRT
|
||||
find_library(NVINFER_LIB nvinfer
|
||||
HINTS /usr/lib /usr/local/lib /opt/tensorrt/lib)
|
||||
find_path(NVINFER_INCLUDE NvInfer.h
|
||||
HINTS /usr/include /usr/local/include /opt/tensorrt/include)
|
||||
if(NOT (NVINFER_LIB AND NVINFER_INCLUDE))
|
||||
message(FATAL_ERROR
|
||||
"TensorRT not found (nvinfer=${NVINFER_LIB} headers=${NVINFER_INCLUDE}). "
|
||||
"Pass -DSAE_INFERENCE_BACKEND=ORT to load .onnx models without TensorRT.")
|
||||
endif()
|
||||
sae_find_cudart()
|
||||
message(STATUS "Inference backend: TRT (${NVINFER_LIB})")
|
||||
|
||||
add_library(inference_backend OBJECT src/backends/trt_backend.cpp)
|
||||
set_target_properties(inference_backend PROPERTIES POSITION_INDEPENDENT_CODE ON)
|
||||
target_include_directories(inference_backend PRIVATE src "${NVINFER_INCLUDE}")
|
||||
target_link_libraries(inference_backend PRIVATE
|
||||
${OpenCV_LIBS} "${NVINFER_LIB}" cudart_dep)
|
||||
target_compile_definitions(inference_backend PRIVATE
|
||||
SAE_MODELS_DIR="${SAE_MODELS_DIR}")
|
||||
endif()
|
||||
|
||||
# ── GEMM backend dependency: builds the `gemm_backend` object lib ──────────────
|
||||
if(SAE_GEMM_BACKEND STREQUAL "CUDA")
|
||||
find_library(CUBLAS_LIB cublas
|
||||
HINTS /opt/cuda/targets/x86_64-linux/lib /opt/cuda/lib64
|
||||
/usr/local/cuda/lib64 /usr/lib)
|
||||
if(NOT CUBLAS_LIB)
|
||||
message(FATAL_ERROR "cuBLAS not found (cublas=${CUBLAS_LIB}).")
|
||||
endif()
|
||||
sae_find_cudart()
|
||||
message(STATUS "GEMM backend: CUDA (${CUBLAS_LIB})")
|
||||
|
||||
add_library(gemm_backend OBJECT src/backends/gemm_backend.cpp)
|
||||
set_target_properties(gemm_backend PROPERTIES POSITION_INDEPENDENT_CODE ON)
|
||||
target_include_directories(gemm_backend PRIVATE src)
|
||||
target_link_libraries(gemm_backend PRIVATE "${CUBLAS_LIB}" cudart_dep)
|
||||
target_compile_definitions(gemm_backend PRIVATE SAE_GEMM_CUDA)
|
||||
else() # ROCM
|
||||
find_library(ROCBLAS_LIB rocblas
|
||||
HINTS /usr/lib64/rocm/lib /usr/lib64 /usr/local/lib)
|
||||
find_path(ROCBLAS_INCLUDE rocblas/rocblas.h
|
||||
HINTS /usr/lib64/rocm/include /usr/include /usr/local/include)
|
||||
find_library(HIP_LIB amdhip64
|
||||
HINTS /usr/lib64/rocm/lib /usr/lib64 /usr/local/lib)
|
||||
find_path(HIP_INCLUDE hip/hip_runtime_api.h
|
||||
HINTS /usr/lib64/rocm/include /usr/include /usr/local/include)
|
||||
if(NOT (ROCBLAS_LIB AND ROCBLAS_INCLUDE AND HIP_LIB AND HIP_INCLUDE))
|
||||
message(FATAL_ERROR
|
||||
"rocBLAS or HIP runtime not found "
|
||||
"(rocblas=${ROCBLAS_LIB} headers=${ROCBLAS_INCLUDE} "
|
||||
"hip=${HIP_LIB} headers=${HIP_INCLUDE}). "
|
||||
"Install rocblas-devel and hip-devel (or pass -DSAE_GEMM_BACKEND=CUDA).")
|
||||
endif()
|
||||
message(STATUS "GEMM backend: ROCM (${ROCBLAS_LIB})")
|
||||
|
||||
add_library(gemm_backend OBJECT src/backends/gemm_backend.cpp)
|
||||
set_target_properties(gemm_backend PROPERTIES POSITION_INDEPENDENT_CODE ON)
|
||||
target_include_directories(gemm_backend PRIVATE src "${ROCBLAS_INCLUDE}" "${HIP_INCLUDE}")
|
||||
target_link_libraries(gemm_backend PRIVATE "${ROCBLAS_LIB}" "${HIP_LIB}")
|
||||
# HIP headers require the platform to be declared explicitly when compiled with g++.
|
||||
target_compile_definitions(gemm_backend PRIVATE SAE_GEMM_ROCM __HIP_PLATFORM_AMD__)
|
||||
endif()
|
||||
|
||||
# FFmpeg (hwaccel video decode: CUDA/VAAPI, runtime-detected + swscale colour
|
||||
# conversion). Hwaccel support is built into libavcodec/libavutil; no extra
|
||||
# libraries are needed here.
|
||||
find_package(PkgConfig REQUIRED)
|
||||
pkg_check_modules(AVFORMAT REQUIRED libavformat)
|
||||
pkg_check_modules(AVCODEC REQUIRED libavcodec)
|
||||
@@ -103,18 +219,22 @@ FetchContent_MakeAvailable(nanobind)
|
||||
set(SAE_MODELS_DIR "${CMAKE_SOURCE_DIR}/models"
|
||||
CACHE PATH "Directory containing ONNX model files")
|
||||
|
||||
# ── Shared library: gallery store ─────────────────────────────────────────────
|
||||
# ── Shared library: gallery store + compiled-in backends ──────────────────────
|
||||
# The backend object libraries carry their own ORT/TRT/CUDA/ROCm linkage and
|
||||
# headers; sae_gallery re-exports those object files so every binary that links
|
||||
# sae_gallery gets the chosen backend without ever seeing its headers.
|
||||
add_library(sae_gallery STATIC
|
||||
src/gallery/gallery_store.cpp
|
||||
src/gallery/gallery_builder.cpp
|
||||
)
|
||||
set_target_properties(sae_gallery PROPERTIES POSITION_INDEPENDENT_CODE ON)
|
||||
target_include_directories(sae_gallery PUBLIC src)
|
||||
target_link_libraries(sae_gallery PUBLIC
|
||||
kpn
|
||||
${OpenCV_LIBS}
|
||||
nlohmann_json::nlohmann_json
|
||||
onnxruntime
|
||||
trt_runtime
|
||||
inference_backend
|
||||
gemm_backend
|
||||
ffmpeg_libs
|
||||
)
|
||||
target_compile_definitions(sae_gallery PUBLIC
|
||||
@@ -123,24 +243,11 @@ target_compile_definitions(sae_gallery PUBLIC
|
||||
|
||||
# ── embed_faces — image → embedding JSON (used by gallery builder scripts) ────
|
||||
add_executable(embed_faces src/embed_faces.cpp)
|
||||
target_link_libraries(embed_faces PRIVATE
|
||||
kpn
|
||||
${OpenCV_LIBS}
|
||||
nlohmann_json::nlohmann_json
|
||||
onnxruntime
|
||||
trt_runtime
|
||||
ffmpeg_libs
|
||||
)
|
||||
target_compile_definitions(embed_faces PRIVATE SAE_MODELS_DIR="${SAE_MODELS_DIR}")
|
||||
target_link_libraries(embed_faces PRIVATE sae_gallery)
|
||||
|
||||
# ── sae_embed — Python module: load SCRFD+ArcFace once, embed many images ───
|
||||
nanobind_add_module(sae_embed src/python_bindings.cpp)
|
||||
target_include_directories(sae_embed PRIVATE src)
|
||||
target_link_libraries(sae_embed PRIVATE
|
||||
${OpenCV_LIBS}
|
||||
onnxruntime
|
||||
)
|
||||
target_compile_definitions(sae_embed PRIVATE SAE_MODELS_DIR="${SAE_MODELS_DIR}")
|
||||
target_link_libraries(sae_embed PRIVATE sae_gallery)
|
||||
|
||||
# ── analyze — main analysis binary ───────────────────────────────────────────
|
||||
add_executable(scene_analyze src/main.cpp)
|
||||
|
||||
Reference in New Issue
Block a user