Skip to content

GitLab

Unverified Commit a59cd9d9 authored by Wentao Ye's avatar Wentao Ye Committed by GitHub
Browse files

[Refactor] Fix Compile Warning #1444-D (#21462) 


Signed-off-by: default avataryewentao256 <zhyanwentao@126.com>
parent 5c54d975
Hide whitespace changes
Inline Side-by-side
Showing with 4 additions and 2 deletions
csrc/moe/topk_softmax_kernels.cu
View file @ a59cd9d9
......@@ -24,9 +24,12 @@
#ifndef USE_ROCM
#include <cub/util_type.cuh>
#include <cub/cub.cuh>
#include <cuda/std/functional>
using AddOp = cuda::std::plus<float>;
#else
#include <hipcub/util_type.hpp>
#include <hipcub/hipcub.hpp>
using AddOp = cub::Sum;
#endif
#define MAX(a, b) ((a) > (b) ? (a) : (b))
......@@ -62,7 +65,6 @@ __launch_bounds__(TPB) __global__
const int thread_row_offset = blockIdx.x * num_cols;
cub::Sum sum;
float threadData(-FLT_MAX);
// Don't touch finished rows.
......@@ -92,7 +94,7 @@ __launch_bounds__(TPB) __global__
threadData += exp((static_cast<float>(input[idx]) - float_max));
}
const auto Z = BlockReduce(tmpStorage).Reduce(threadData, sum);
const auto Z = BlockReduce(tmpStorage).Reduce(threadData, AddOp());
if (threadIdx.x == 0)
{
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment