From 51747d90c31d13116fdcbb4ca188ac4c03f30389 Mon Sep 17 00:00:00 2001
From: Muhammed Fatih BALIN <m.f.balin@gmail.com>
Date: Tue, 5 Dec 2023 02:22:36 -0500
Subject: [PATCH] [Bug] Update thrust to the new monorepo CCCL (#6166)

Co-authored-by: Rhett Ying <85214957+Rhett-Ying@users.noreply.github.com>
---
 .gitmodules                             | 6 +++---
 CMakeLists.txt                          | 8 ++++----
 graphbolt/CMakeLists.txt                | 8 ++++----
 graphbolt/src/cuda/index_select_impl.cu | 2 +-
 third_party/cccl                        | 1 +
 third_party/thrust                      | 1 -
 6 files changed, 13 insertions(+), 13 deletions(-)
 create mode 160000 third_party/cccl
 delete mode 160000 third_party/thrust

diff --git a/.gitmodules b/.gitmodules
index c35dd0013be4..9b45e6f5a305 100644
--- a/.gitmodules
+++ b/.gitmodules
@@ -19,9 +19,9 @@
 [submodule "third_party/libxsmm"]
 	path = third_party/libxsmm
 	url = https://github.com/hfp/libxsmm.git
-[submodule "third_party/thrust"]
-	path = third_party/thrust
-	url = https://github.com/NVIDIA/thrust.git
 [submodule "third_party/pcg"]
 	path = third_party/pcg
 	url = https://github.com/imneme/pcg-cpp.git
+[submodule "third_party/cccl"]
+	path = third_party/cccl
+	url = https://github.com/NVIDIA/cccl.git
diff --git a/CMakeLists.txt b/CMakeLists.txt
index afbe6d421c98..ea5d9267db5a 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -142,10 +142,10 @@ if(USE_CUDA)
   # see https://github.com/NVIDIA/thrust/issues/1401
   add_definitions(-DTHRUST_CUB_WRAPPED_NAMESPACE=dgl)
   include(cmake/modules/CUDA.cmake)
-  message(STATUS "Use external CUB/Thrust library for a consistent API and performance.")
-  cuda_include_directories(BEFORE "${CMAKE_SOURCE_DIR}/third_party/thrust")
-  cuda_include_directories(BEFORE "${CMAKE_SOURCE_DIR}/third_party/thrust/dependencies/cub")
-  cuda_include_directories(BEFORE "${CMAKE_SOURCE_DIR}/third_party/thrust/dependencies/libcudacxx/include")
+  message(STATUS "Use external CCCL library for a consistent API and performance.")
+  cuda_include_directories(BEFORE "${CMAKE_SOURCE_DIR}/third_party/cccl/thrust")
+  cuda_include_directories(BEFORE "${CMAKE_SOURCE_DIR}/third_party/cccl/cub")
+  cuda_include_directories(BEFORE "${CMAKE_SOURCE_DIR}/third_party/cccl/libcudacxx/include")
 endif(USE_CUDA)
 
 # initial variables
diff --git a/graphbolt/CMakeLists.txt b/graphbolt/CMakeLists.txt
index 14511385ffc7..85a238eeed5d 100644
--- a/graphbolt/CMakeLists.txt
+++ b/graphbolt/CMakeLists.txt
@@ -60,13 +60,13 @@ target_link_libraries(${LIB_GRAPHBOLT_NAME} "${TORCH_LIBRARIES}")
 
 if(USE_CUDA)
   set_target_properties(${LIB_GRAPHBOLT_NAME} PROPERTIES CUDA_STANDARD 17)
-  message(STATUS "Use external CUB/Thrust library for a consistent API and performance for graphbolt.")
+  message(STATUS "Use external CCCL library for a consistent API and performance for graphbolt.")
   target_compile_definitions(${LIB_GRAPHBOLT_NAME} PRIVATE CUB_WRAPPED_NAMESPACE=graphbolt)
   target_compile_definitions(${LIB_GRAPHBOLT_NAME} PRIVATE THRUST_NS_QUALIFIER=thrust)
   target_include_directories(${LIB_GRAPHBOLT_NAME} PRIVATE
-                             "../third_party/thrust"
-                             "../third_party/thrust/dependencies/cub"
-                             "../third_party/thrust/dependencies/libcudacxx/include")
+                             "../third_party/cccl/thrust"
+                             "../third_party/cccl/cub"
+                             "../third_party/cccl/libcudacxx/include")
 endif()
 
 # The Torch CMake configuration only sets up the path for the MKL library when
diff --git a/graphbolt/src/cuda/index_select_impl.cu b/graphbolt/src/cuda/index_select_impl.cu
index dd07e1144c6a..47b04f439d82 100644
--- a/graphbolt/src/cuda/index_select_impl.cu
+++ b/graphbolt/src/cuda/index_select_impl.cu
@@ -6,12 +6,12 @@
 #include <c10/core/ScalarType.h>
 #include <c10/cuda/CUDAStream.h>
 
+#include <cub/cub.cuh>
 #include <numeric>
 
 #include "../index_select.h"
 #include "./common.h"
 #include "./utils.h"
-#include "cub/cub.cuh"
 
 namespace graphbolt {
 namespace ops {
diff --git a/third_party/cccl b/third_party/cccl
new file mode 160000
index 000000000000..4d5c181cb4e6
--- /dev/null
+++ b/third_party/cccl
@@ -0,0 +1 @@
+Subproject commit 4d5c181cb4e67b14e7d1fdeb45cf0104119bfb09
diff --git a/third_party/thrust b/third_party/thrust
deleted file mode 160000
index 02931a309bee..000000000000
--- a/third_party/thrust
+++ /dev/null
@@ -1 +0,0 @@
-Subproject commit 02931a309bee769853088b79b4e3ab1c0bd2336c