update

.gitmodules

@@ -7,3 +7,9 @@
 [submodule "third_party/METIS"]
     path = third_party/METIS
     url = https://github.com/KarypisLab/METIS.git
+[submodule "third_party/cuCollections"]
+    path = third_party/cuCollections
+    url = https://github.com/NVIDIA/cuCollections.git
+[submodule "third_party/cccl"]
+    path = third_party/cccl
+    url = https://github.com/NVIDIA/cccl.git

CMakeLists.txt

@@ -52,6 +52,8 @@ if(WITH_CUDA)
     set(CUTLASS_UTIL_DIR third_party/cutlass/tools/util/include)
     include_directories(${CUTLASS_UTIL_DIR})
   endif()
+  set(CUCOLLECTIONS_DIR third_party/cuCollections/include)
+  include_directories(${CUCOLLECTIONS_DIR})
 endif()
 
 set(CSRC pyg_lib/csrc)
@@ -105,6 +107,13 @@ if (USE_PYTHON)
   target_link_libraries(${PROJECT_NAME} PRIVATE Python3::Python)
 endif()
 
+if(WITH_CUDA)
+  target_include_directories(${PROJECT_NAME} PRIVATE
+                             third_party/cccl/thrust
+                             third_party/cccl/cub
+                             third_party/cccl/libcudacxx/include)
+endif()
+
 set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -O0")
 
 if(BUILD_TEST)

pyg_lib/csrc/ops/cuda/sampled_kernel.cu

@@ -2,13 +2,15 @@
 #include <ATen/cuda/CUDAContext.h>
 #include <torch/library.h>
 
+#include <cuco/static_map.cuh>
+
 namespace pyg {
 namespace ops {
 
 namespace {
 
 #define THREADS 1024
-#define CDIV(N, M) ((N) + (M) - 1) / (M)
+#define CDIV(N, M) ((N) + (M)-1) / (M)
 
 enum FnType { ADD, SUB, MUL, DIV };
 const std::map<std::string, FnType> to_fn_type = {