Fix building for CUDA 11.6 (#917)

* Fix building for CUDA 11.6

* Minor fixes.

Author: csukuangfj

k2/csrc/CMakeLists.txt

@@ -92,6 +92,11 @@ add_library(context ${context_srcs})
 target_compile_definitions(context PUBLIC K2_TORCH_VERSION_MAJOR=${K2_TORCH_VERSION_MAJOR})
 target_compile_definitions(context PUBLIC K2_TORCH_VERSION_MINOR=${K2_TORCH_VERSION_MINOR})
 
+# see https://github.com/NVIDIA/thrust/issues/1401
+# and https://github.com/k2-fsa/k2/pull/917
+target_compile_definitions(context PUBLIC CUB_WRAPPED_NAMESPACE=k2)
+target_compile_definitions(context PUBLIC THRUST_NS_QUALIFIER=thrust)
+
 set_target_properties(context PROPERTIES CUDA_SEPARABLE_COMPILATION ON)
 set_target_properties(context PROPERTIES OUTPUT_NAME "k2context")
 

k2/csrc/cub.h

@@ -20,50 +20,8 @@
 #ifndef K2_CSRC_CUB_H_
 #define K2_CSRC_CUB_H_
 
-// See
-// https://github.com/k2-fsa/k2/issues/698
-// and
-// https://github.com/pytorch/pytorch/issues/54245#issuecomment-805707551
-// for why we need the following two macros
-//
-// NOTE: We define the following two macros so
-// that k2 and PyTorch use a different copy
-// of CUB.
-
-#ifdef CUB_NS_PREFIX
-#undef CUB_NS_PREFIX
-#endif
-
-#ifdef CUB_NS_POSTFIX
-#undef CUB_NS_POSTFIX
-#endif
-
-#ifdef CUB_NS_QUALIFIER
-#undef CUB_NS_QUALIFIER
-#endif
-
-// see
-// https://github.com/NVIDIA/cub/commit/6631c72630f10e370d93814a59146b12f7620d85
-// The above commit replaced "thrust" with "THRUST_NS_QUALIFIER"
-#ifndef THRUST_NS_QUALIFIER
-#define THRUST_NS_QUALIFIER thrust
-#endif
-
-#define CUB_NS_PREFIX namespace k2 {
-#define CUB_NS_POSTFIX }
-
-// See
-// https://github.com/NVIDIA/cub/commit/6631c72630f10e370d93814a59146b12f7620d85
-// and
-// https://github.com/NVIDIA/cub/pull/350
-#define CUB_NS_QUALIFIER ::k2::cub
-
 #ifdef K2_WITH_CUDA
 #include "cub/cub.cuh"  // NOLINT
 #endif
 
-#undef CUB_NS_PREFIX
-#undef CUB_NS_POSTFIX
-#undef CUB_NS_QUALIFIER
-
 #endif  // K2_CSRC_CUB_H_