[BugFix]Fix cub reduce bug in cuda12 (#58179)

YuanRisheng · web-flow · commit 6c8178bad752 · 2023-11-22T14:38:30.000+08:00
* open shared phi

* fix inference bugs

* fix py3 bugs

* hot fix

* fix

* close phi shared

* open cub
diff --git a/cmake/external/cccl.cmake b/cmake/external/cccl.cmake
@@ -26,6 +26,7 @@ ExternalProject_Add(
   INSTALL_COMMAND ""
   TEST_COMMAND "")
 
-add_library(cccl INTERFACE)
-
-add_dependencies(cccl extern_cccl)
+# update include dir and set cccl first for using
+include_directories(BEFORE "${CCCL_SOURCE_DIR}/cub")
+include_directories(BEFORE "${CCCL_SOURCE_DIR}/libcudacxx/include")
+include_directories(BEFORE "${CCCL_SOURCE_DIR}/thrust")
diff --git a/cmake/generic.cmake b/cmake/generic.cmake
@@ -1348,9 +1348,6 @@ function(math_library TARGET)
   if(WITH_GPU)
     if(${CMAKE_CUDA_COMPILER_VERSION} LESS 11.0)
       list(APPEND math_common_deps cub)
-    elseif(${CMAKE_CUDA_COMPILER_VERSION} EQUAL 12.0
-           OR ${CMAKE_CUDA_COMPILER_VERSION} GREATER 12.0)
-      list(APPEND math_common_deps cccl)
     else()
       list(APPEND math_common_deps)
     endif()
diff --git a/cmake/third_party.cmake b/cmake/third_party.cmake
@@ -397,7 +397,7 @@ if(WITH_GPU)
   elseif(${CMAKE_CUDA_COMPILER_VERSION} EQUAL 12.0
          OR ${CMAKE_CUDA_COMPILER_VERSION} GREATER 12.0)
     include(external/cccl)
-    list(APPEND third_party_deps extern_cccl)
+    add_definitions(-DPADDLE_WITH_CCCL)
   endif()
   set(URL
       "https://paddlepaddledeps.bj.bcebos.com/externalErrorMsg_20210928.tar.gz"
diff --git a/paddle/phi/common/complex.h b/paddle/phi/common/complex.h
@@ -74,6 +74,14 @@ struct PADDLE_ALIGN(sizeof(T) * 2) complex {
     imag = c.imag();
   }
 
+#if defined(PADDLE_WITH_CCCL)
+  template <typename T1>
+  HOSTDEVICE inline explicit complex(const cuda::std::complex<T1>& c) {
+    real = c.real();
+    imag = c.imag();
+  }
+#endif
+
   template <typename T1>
   HOSTDEVICE inline explicit operator thrust::complex<T1>() const {
     return thrust::complex<T1>(real, imag);
diff --git a/paddle/phi/kernels/funcs/reduce_function.h b/paddle/phi/kernels/funcs/reduce_function.h
@@ -1096,16 +1096,6 @@ void ReduceKernel(const KPDevice& dev_ctx,
   constexpr bool kIsTxBF16 = std::is_same<Tx, phi::dtype::bfloat16>::value;
   bool use_cub_reduce = config.reduce_num == numel && !kIsTxFP16 && !kIsTxBF16;
 
-  // NOTE(YuanRisheng): hot fix
-  // cuda 12.0 + cub got wrong result in some shapes when build phi with shared
-  // library. For example, paddle.sum(paddle.ones([1024,100],
-  // dtype=paddle.float32)) is expected to 102400, but got 0.
-#ifdef PHI_SHARED
-#if CUDA_VERSION >= 12000
-  use_cub_reduce = false;
-#endif
-#endif
-
 #ifndef PADDLE_WITH_XPU_KP
   if (use_cub_reduce) {
     CubTensorReduce<Tx, Ty, ReduceOp, TransformOp, IsMean>::apply(
diff --git a/paddle/phi/kernels/fusion/gpu/fused_scale_bias_relu_conv_bn_kernel.cu b/paddle/phi/kernels/fusion/gpu/fused_scale_bias_relu_conv_bn_kernel.cu
@@ -12,6 +12,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License. */
 
+#include <float.h>
 #include <array>
 
 #include "paddle/phi/backends/gpu/cuda/cudnn_helper.h"
