From 329306b031d10ca99756e1baf1e090bfc2ae73eb Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Mon, 25 Jan 2021 14:55:07 +0800
Subject: [PATCH] fix(cmake/cuda): fix build at cuda `copy` env caused by
 b278a69e1 `copy` env vs `normal` env: case a:`normal` env install CUDNN/TRT
 by package tool, eg, apt/yum, which will put include file(cudnn.h/NvInfer.h)
 to /usr/include etc.

case b: self copy env after export CPATH with cudnn/trt include dir

case c: just copy libs, not export CPATH, only export
TRT_ROOT_DIR/CUDNN_ROOT_DIR

at case a and case b, BUILD system can find cudnn/trt headers, but
at case c, can not find. `FIX` it!

what`s more, recommend install CUDA env by case a and case c, NOT
recommend install by case b, which caused by CUDA SDK CPATH not
friendly to split c env, for example Android AOSP build env

GitOrigin-RevId: cb92123f9e58d730893fff45602d3e487b38d32c
---
 cmake/cudnn.cmake      | 4 ++++
 cmake/tensorrt.cmake   | 4 ++++
 dnn/src/CMakeLists.txt | 1 +
 src/CMakeLists.txt     | 3 +++
 4 files changed, 12 insertions(+)

diff --git a/cmake/cudnn.cmake b/cmake/cudnn.cmake
index 6011e3e9..3fd42803 100644
--- a/cmake/cudnn.cmake
+++ b/cmake/cudnn.cmake
@@ -7,6 +7,10 @@ if(NOT "$ENV{LIBRARY_PATH}" STREQUAL "")
     string(REPLACE ":" ";" SYSTEM_LIBRARY_PATHS $ENV{LIBRARY_PATH})
 endif()
 
+if("${CUDNN_ROOT_DIR}" STREQUAL "" AND NOT "$ENV{CUDNN_ROOT_DIR}"  STREQUAL "")
+    set(CUDNN_ROOT_DIR $ENV{CUDNN_ROOT_DIR})
+endif()
+
 if(MGE_CUDA_USE_STATIC)
     find_library(CUDNN_LIBRARY 
         NAMES libcudnn_static.a cudnn.lib
diff --git a/cmake/tensorrt.cmake b/cmake/tensorrt.cmake
index 25d21aea..7ac5b5b2 100644
--- a/cmake/tensorrt.cmake
+++ b/cmake/tensorrt.cmake
@@ -2,6 +2,10 @@ if(NOT "$ENV{LIBRARY_PATH}" STREQUAL "")
     string(REPLACE ":" ";" SYSTEM_LIBRARY_PATHS $ENV{LIBRARY_PATH})
 endif()
 
+if("${TRT_ROOT_DIR}" STREQUAL "" AND NOT "$ENV{TRT_ROOT_DIR}"  STREQUAL "")
+    set(TRT_ROOT_DIR $ENV{TRT_ROOT_DIR})
+endif()
+
 if(MGE_CUDA_USE_STATIC)
     find_library(TRT_LIBRARY 
         NAMES libnvinfer_static.a nvinfer.lib
diff --git a/dnn/src/CMakeLists.txt b/dnn/src/CMakeLists.txt
index 430a6b94..029f1bf8 100644
--- a/dnn/src/CMakeLists.txt
+++ b/dnn/src/CMakeLists.txt
@@ -134,6 +134,7 @@ add_library(megdnn EXCLUDE_FROM_ALL OBJECT ${SOURCES})
 target_link_libraries(megdnn PUBLIC opr_param_defs)
 if(MGE_WITH_CUDA)
     target_link_libraries(megdnn PRIVATE $<BUILD_INTERFACE:cutlass>)
+    target_include_directories(megdnn PRIVATE ${CUDNN_INCLUDE_DIR})
 endif()
 
 if(MGE_WITH_ROCM)
diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
index 8c0db789..5fddd520 100644
--- a/src/CMakeLists.txt
+++ b/src/CMakeLists.txt
@@ -57,6 +57,9 @@ endif()
 
 add_library(megbrain OBJECT ${SOURCES})
 target_link_libraries(megbrain PUBLIC mgb_opr_param_defs)
+if(MGE_WITH_CUDA)
+    target_include_directories(megbrain PUBLIC ${TRT_INCLUDE_DIR})
+endif()
 target_include_directories(megbrain
     PUBLIC $<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>
     PRIVATE ${PROJECT_SOURCE_DIR}/third_party/midout/src
-- 
GitLab