fix BlockPrefixCallbackOp

Author: cangtianhuang

paddle/phi/kernels/gpu/cum_kernel.cu

@@ -128,18 +128,36 @@ struct Identity<T, ComplexSum> {
  static constexpr T value = {0, 0};
 };
 
+template <typename T, typename Op, bool UseKahan>
+struct BlockPrefixCallbackOp;
+
 template <typename T, typename Op>
-struct BlockPrefixCallbackOp {
+struct BlockPrefixCallbackOp<T, Op, false> {
  // Running prefix
  T running_total_;
- T compensation_;
  Op op_;
 
  __device__ BlockPrefixCallbackOp(T identity, Op op)
- : running_total_(identity), compensation_(identity), op_(op) {}
+ : running_total_(identity), op_(op) {}
 
  // Callback operator to be entered by the first warp of threads in the block.
  // tid 0 is responsible for returning a value for seeding the block-wide scan.
+ __device__ T operator()(T block_aggregate) {
+ const T old_prefix = running_total_;
+ running_total_ = op_(running_total_, block_aggregate);
+ return old_prefix;
+ }
+};
+
+template <typename T, typename Op>
+struct BlockPrefixCallbackOp<T, Op, true> {
+ T running_total_;
+ T compensation_;
+ Op op_;
+
+ __device__ BlockPrefixCallbackOp(T identity, Op op)
+ : running_total_(identity), compensation_(static_cast<T>(0.0)), op_(op) {}
+
  __device__ T operator()(T block_aggregate) {
  T old_prefix = running_total_;
 
@@ -155,20 +173,23 @@ struct BlockPrefixCallbackOp {
 };
 
 template <typename T>
-struct BlockPrefixCallbackOp<T, LogAddExp> {
+struct BlockPrefixCallbackOp<T, LogAddExp, true> {
  T max_so_far_;
  T scaled_sum_;
  T compensation_;
  LogAddExp op_;
 
  __device__ BlockPrefixCallbackOp(T identity, LogAddExp op)
- : max_so_far_(identity), scaled_sum_(0.0), compensation_(0.0), op_(op) {}
+ : max_so_far_(identity),
+ scaled_sum_(static_cast<T>(0.0)),
+ compensation_(static_cast<T>(0.0)),
+ op_(op) {}
 
  __device__ T operator()(T block_aggregate) {
  if (scaled_sum_ == 0.0) {
  max_so_far_ = block_aggregate;
- scaled_sum_ = 1.0;
- compensation_ = 0.0;
+ scaled_sum_ = static_cast<T>(1.0);
+ compensation_ = static_cast<T>(0.0);
  return std::numeric_limits<T>::lowest();
  }
 
@@ -195,15 +216,19 @@ struct BlockPrefixCallbackOp<T, LogAddExp> {
  }
 };
 
-template <typename T, int BLOCK_THREADS, int ITEMS_PER_THREAD, typename Op>
+template <typename T,
+ int BLOCK_THREADS,
+ int ITEMS_PER_THREAD,
+ typename Op,
+ bool UseKahan>
 __global__ void BlockScanKernel(T* d_out,
  const T* d_in,
  int64_t grid_size,
  int64_t scan_size,
  bool exclusive,
  Op op) {
  using MT = typename phi::dtype::MPTypeTrait<T>::Type;
- using CallbackOp = BlockPrefixCallbackOp<MT, Op>;
+ using CallbackOp = BlockPrefixCallbackOp<MT, Op, UseKahan>;
 
  // Specialize BlockLoad, BlockStore, and BlockRadixSort collective types
  using BlockLoadT = cub::
@@ -350,14 +375,30 @@ void ScanKernel(const Context& dev_ctx,
  }
  }
 
+ // When scan_size is large, switch to Kahan scan to get better precision
+ constexpr int64_t KAHAN_SWITCH_LENGTH = 1 << 16;
+
  // Do scan
  if (!transpose && !reverse) {
- BlockScanKernel<T, 128, 4, Op><<<scan_grid, 128, 0, dev_ctx.stream()>>>(
- out_data, in_data, grid_size, scan_size, exclusive, op);
-
+ if (scan_size > KAHAN_SWITCH_LENGTH) {
+ BlockScanKernel<T, 128, 4, Op, true>
+ <<<scan_grid, 128, 0, dev_ctx.stream()>>>(
+ out_data, in_data, grid_size, scan_size, exclusive, op);
+ } else {
+ BlockScanKernel<T, 128, 4, Op, false>
+ <<<scan_grid, 128, 0, dev_ctx.stream()>>>(
+ out_data, in_data, grid_size, scan_size, exclusive, op);
+ }
  } else {
- BlockScanKernel<T, 128, 4, Op><<<scan_grid, 128, 0, dev_ctx.stream()>>>(
- next_out_data, next_in_data, grid_size, scan_size, exclusive, op);
+ if (scan_size > KAHAN_SWITCH_LENGTH) {
+ BlockScanKernel<T, 128, 4, Op, true>
+ <<<scan_grid, 128, 0, dev_ctx.stream()>>>(
+ next_out_data, next_in_data, grid_size, scan_size, exclusive, op);
+ } else {
+ BlockScanKernel<T, 128, 4, Op, false>
+ <<<scan_grid, 128, 0, dev_ctx.stream()>>>(
+ next_out_data, next_in_data, grid_size, scan_size, exclusive, op);
+ }
  }
  swap_ptr(next_in_data, next_out_data);