From 83520fd232610facf743a0f9728d3749e8a707bd Mon Sep 17 00:00:00 2001
From: Sing_chan <51314274+betterpig@users.noreply.github.com>
Date: Thu, 30 Jun 2022 14:53:09 +0800
Subject: [PATCH] cherry pick 43934 and not format (#43935)

---
 cmake/external/cub.cmake | 10 +++++++++-
 cmake/third_party.cmake  |  3 ++-
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/cmake/external/cub.cmake b/cmake/external/cub.cmake
index f263086e8be..8f386823ff9 100644
--- a/cmake/external/cub.cmake
+++ b/cmake/external/cub.cmake
@@ -21,7 +21,15 @@ set(CUB_PATH        "${THIRD_PARTY_PATH}/cub" CACHE STRING "A path setting for e
 set(CUB_PREFIX_DIR  ${CUB_PATH})
 
 set(CUB_REPOSITORY  ${GIT_URL}/NVlabs/cub.git)
-set(CUB_TAG         1.8.0)
+
+if(WIN32 AND ${CMAKE_CUDA_COMPILER_VERSION} GREATER_EQUAL 11.6)
+  # cuda_11.6.2_511.65‘s own cub is 1.15.0, which will cause compiling error in windows.
+  set(CUB_TAG 1.16.0)
+  # cub 1.16.0 is not compitable with current thrust version
+  add_definitions(-DTHRUST_IGNORE_CUB_VERSION_CHECK)
+else()
+  set(CUB_TAG 1.8.0)
+endif()
 
 SET(CUB_INCLUDE_DIR  ${CUB_PREFIX_DIR}/src/extern_cub)
 message("CUB_INCLUDE_DIR is ${CUB_INCLUDE_DIR}")
diff --git a/cmake/third_party.cmake b/cmake/third_party.cmake
index 7df095c6c2e..3c39d508789 100644
--- a/cmake/third_party.cmake
+++ b/cmake/third_party.cmake
@@ -257,7 +257,8 @@ if(WITH_ONNXRUNTIME)
 endif()
 
 if(WITH_GPU)
-    if (${CMAKE_CUDA_COMPILER_VERSION} LESS 11.0)
+    if (${CMAKE_CUDA_COMPILER_VERSION} LESS 11.0 OR ${CMAKE_CUDA_COMPILER_VERSION}
+                                                    GREATER_EQUAL 11.6)
         include(external/cub)       # download cub
         list(APPEND third_party_deps extern_cub)
     endif()
-- 
GitLab