Drop NVRTC_CUB in non-RTC paths

There is no need to use modified CUB in nvcc or LLVM paths because,
unlike NVRTC, they can include system headers that the regular CUB
depends upon.

Author: nicolasvasilache

tc/core/cuda/cuda_libraries.h

@@ -31,9 +31,11 @@ namespace code {
 namespace c {
 
 constexpr auto types = R"C(
+#ifndef __CUDACC_RTC__
 // Can't include system dependencies with NVRTC
 // Can't include cuda_fp16.h with NVRTC due to transitive system dependencies
-// #include <cuda_fp16.h>
+#include <cuda_fp16.h>
+#endif
 )C";
 
 constexpr auto defines = R"C(
@@ -211,7 +213,12 @@ struct SegmentedReducer {
 
 constexpr auto cubBlockReduce = R"CUDA(
 
+#if __CUDACC_RTC__
 #include "cub/nvrtc_cub.cuh"
+#else
+#include <assert.h>
+#include "cub/cub.cuh"
+#endif
 
 namespace __tc {
 

tc/core/cuda/cuda_rtc.cc

@@ -128,7 +128,6 @@ static std::string llvmCompile(
        std::string("-I") + TC_STRINGIFY(TC_CUDA_INCLUDE_DIR),
        std::string("-I") + TC_STRINGIFY(TC_CUB_INCLUDE_DIR),
        tc::FLAGS_llvm_flags,
-       "-DNVRTC_CUB=1",
        "-nocudalib",
        "-S",
        "-emit-llvm",

tc/core/flags.cc

@@ -49,7 +49,7 @@ DEFINE_string(
     "compiler flags to set when llvm is used");
 DEFINE_string(
     nvcc_flags,
-    "-std=c++11 -ptx -DNVRTC_CUB=1 --use_fast_math",
+    "-std=c++11 -ptx --use_fast_math",
     "compiler flags to set when nvcc is used");
 
 // CPU codegen options