saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-26 19:52:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Bugfix for topk cuda kernel (#6164)

Browse source 
* fix the issue that std::numeric_limits cannot handle half type

* adding a test

Co-authored-by: Du Li <duli@OrtTrainingDev4.af05slrtruoetgaxwwjv5nsq5e.px.internal.cloudapp.net>
This commit is contained in:
Du Li 2020-12-17 17:59:37 -08:00 committed by GitHub
parent dec703b62d
commit 34725ae520
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 46 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

37
onnxruntime/core/providers/cuda/math/topk_impl.cu
View file

@ -22,6 +22,26 @@ struct KV {
int64_t val;
};
template <typename T>
struct NumericLimits {
static T Lowest() {
return std::numeric_limits<T>::lowest();
}
static T Max() {
return std::numeric_limits<T>::max();
}
};
template <>
struct NumericLimits<MLFloat16> {
static half Lowest() {
return -65504.0;
}
static half Max() {
return 65504.0;
}
};
#define BT GridDim::maxThreadsPerBlock
#define ALIGN(N) static_cast<int64_t>(pow(2, ceil(log2(static_cast<double>(N)))))
#define FROM(idx) (left_dim + (idx)*mid_dim + right_dim)
@ -147,6 +167,7 @@ __global__ void BitonicTopK(const T* X, T* V, int64_t* I, const TArray<int64_t>
}
}
template <typename T>
__device__ __inline__ bool Equal(const T& t0, const T& t1) {
return t0 == t1;
@ -168,7 +189,7 @@ __device__ bool SamePrefix(const T* t0, const T* t1, int64_t skip) {
}
__device__ bool SamePrefix(const half* f0, const half* f1, int64_t skip) {
return SamePrefix((const uint16_t*)f0, (const uint16_t*)f1, skip);
return SamePrefix((const int16_t*)f0, (const int16_t*)f1, skip);
}
__device__ bool SamePrefix(const float* f0, const float* f1, int64_t skip) {
@ -185,7 +206,7 @@ __device__ int32_t Radix(const T* t, int64_t skip) {
}
__device__ int32_t Radix(const half* f, int64_t skip) {
return Radix((const uint16_t*)f, skip);
return Radix((const int16_t*)f, skip);
}
__device__ int32_t Radix(const float* f, int64_t skip) {
@ -202,7 +223,7 @@ __device__ void SetByte(T* t, int64_t byte) {
}
__device__ void SetByte(half* f, int64_t byte) {
SetByte((uint16_t*)f, byte);
SetByte((int16_t*)f, byte);
}
__device__ void SetByte(float* f, int64_t byte) {
@ -401,17 +422,17 @@ Status TopKImpl(const CudaKernel* kernel, const T* input_x, T* output_v, int64_t
auto aligned_K = ALIGN(K);
auto aligned_dimension = ALIGN(dimension);
if (aligned_dimension <= GridDim::maxThreadsPerBlock) {
BitonicTopK<CudaT><<<N, GridDim::maxThreadsPerBlock, aligned_dimension * sizeof(KV<CudaT>)>>>(input_x_ptr, output_v_ptr, output_i, elem_nums, size, axis, K, aligned_K, largest, sorted, dimension, aligned_dimension, std::numeric_limits<CudaT>::lowest(), std::numeric_limits<CudaT>::max());
BitonicTopK<CudaT><<<N, GridDim::maxThreadsPerBlock, aligned_dimension * sizeof(KV<CudaT>)>>>(input_x_ptr, output_v_ptr, output_i, elem_nums, size, axis, K, aligned_K, largest, sorted, dimension, aligned_dimension, NumericLimits<T>::Lowest(), NumericLimits<T>::Max());
} else if (K <= BT*16 || 0 == sorted) {
auto XPT = static_cast<int64_t>(ceil(static_cast<double>(dimension) / GridDim::maxThreadsPerBlock));
if (BT*2 >= K || 0 == sorted) {
RadixTopK<CudaT,BT,2><<<N,BT,256*sizeof(uint32_t)>>>(input_x_ptr, output_v_ptr, output_i, elem_nums, size, axis, K, largest, sorted, dimension, XPT, std::numeric_limits<CudaT>::lowest(), std::numeric_limits<CudaT>::max());
RadixTopK<CudaT, BT, 2><<<N, BT, 256 * sizeof(uint32_t)>>>(input_x_ptr, output_v_ptr, output_i, elem_nums, size, axis, K, largest, sorted, dimension, XPT, NumericLimits<T>::Lowest(), NumericLimits<T>::Max());
} else if (BT*4>=K) {
RadixTopK<CudaT,BT,4><<<N,BT,256*sizeof(uint32_t)>>>(input_x_ptr, output_v_ptr, output_i, elem_nums, size, axis, K, largest, sorted, dimension, XPT, std::numeric_limits<CudaT>::lowest(), std::numeric_limits<CudaT>::max());
RadixTopK<CudaT, BT, 4><<<N, BT, 256 * sizeof(uint32_t)>>>(input_x_ptr, output_v_ptr, output_i, elem_nums, size, axis, K, largest, sorted, dimension, XPT, NumericLimits<T>::Lowest(), NumericLimits<T>::Max());
} else if (BT*8>=K) {
RadixTopK<CudaT,BT,8><<<N,BT,256*sizeof(uint32_t)>>>(input_x_ptr, output_v_ptr, output_i, elem_nums, size, axis, K, largest, sorted, dimension, XPT, std::numeric_limits<CudaT>::lowest(), std::numeric_limits<CudaT>::max());
RadixTopK<CudaT, BT, 8><<<N, BT, 256 * sizeof(uint32_t)>>>(input_x_ptr, output_v_ptr, output_i, elem_nums, size, axis, K, largest, sorted, dimension, XPT, NumericLimits<T>::Lowest(), NumericLimits<T>::Max());
} else {
RadixTopK<CudaT,BT,16><<<N,BT,256*sizeof(uint32_t)>>>(input_x_ptr, output_v_ptr, output_i, elem_nums, size, axis, K, largest, sorted, dimension, XPT, std::numeric_limits<CudaT>::lowest(), std::numeric_limits<CudaT>::max());
RadixTopK<CudaT, BT, 16><<<N, BT, 256 * sizeof(uint32_t)>>>(input_x_ptr, output_v_ptr, output_i, elem_nums, size, axis, K, largest, sorted, dimension, XPT, NumericLimits<T>::Lowest(), NumericLimits<T>::Max());
}
} else {
auto input_key_buffer = kernel->GetScratchBuffer<CudaT>(dimension);

17
onnxruntime/test/providers/cpu/math/topk_op_test.cc
View file

@ -456,6 +456,23 @@ TEST(TopKOperator, NthElementHalf) {
RunTest(11, 4, input_vals, input_dimensions, expected_vals, expected_indices, expected_dimensions, false);
}
TEST(TopKOperator, NthElementHalf_NegtiveVals) {
if (!HasCudaEnvironment(600)) {
return;
}
std::vector<float> input_vals_f = {10.0f, -8.0f, -7.0f, -4.0f, -5.0f, -6.0f};
std::vector<float> expected_vals_f = {10.0f, -4.0f, -5.0f, -6.0f};
std::vector<MLFloat16> input_vals(6);
std::vector<MLFloat16> expected_vals(4);
ConvertFloatToMLFloat16(input_vals_f.data(), input_vals.data(), 6);
ConvertFloatToMLFloat16(expected_vals_f.data(), expected_vals.data(), 4);
std::vector<int64_t> input_dimensions = {6};
std::vector<int64_t> expected_indices = {0, 3, 4, 5};
std::vector<int64_t> expected_dimensions = {4};
RunTest(11, 4, input_vals, input_dimensions, expected_vals, expected_indices, expected_dimensions, false);
}
// test dimension in range (GridDim::maxThreadsPerBlock, GridDim::maxThreadsPerBlock * 2], ie. [257, 512]
TEST(TopKOperator, SmallArrayTopKSorted) {
std::vector<float> input_vals(400, 0.0f);