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

Fix TopK Cuda implementation (#3176)

Browse source 
Fixes a bug in TopK cuda implementation when input size is between GridDim::maxThreadsPerBlock and GridDim::maxThreadsPerBlock * 2. In this case, the BitonicTopK will generate all-zero outputs.
This commit is contained in:
Yulong Wang 2020-03-13 11:46:17 -07:00 committed by GitHub
parent 93569bf0f4
commit 5bc0d8be5c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 19 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -372,7 +372,7 @@ template <typename T>
Status TopKImpl(const CudaKernel* kernel, const T* input_x, T* output_v, int64_t* output_i, const int64_t* elem_nums, size_t size, int64_t axis, int64_t K, int64_t largest, int64_t sorted, int64_t N, int64_t dimension) {
auto aligned_K = ALIGN(K);
auto aligned_dimension = ALIGN(dimension);
if (aligned_dimension <= GridDim::maxThreadsPerBlock << 1) {
if (aligned_dimension <= GridDim::maxThreadsPerBlock) {
BitonicTopK<T><<<N, GridDim::maxThreadsPerBlock, aligned_dimension * sizeof(KV<T>)>>>(input_x, output_v, output_i, elem_nums, size, axis, K, aligned_K, largest, sorted, dimension, aligned_dimension, std::numeric_limits<T>::lowest(), std::numeric_limits<T>::max());
} else if (K <= BT*16 || 0 == sorted) {
auto XPT = static_cast<int64_t>(ceil(static_cast<double>(dimension) / GridDim::maxThreadsPerBlock));

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

@ -475,8 +475,22 @@ TEST(TopKOperator, SortedSelection) {
RunTest(11, 5, input_vals, input_dimensions, expected_vals, expected_indices, expected_dimensions, false, axis, 0); // smallest values
}
TEST(TopKOperator, MediumArrayTopKSorted)
{
// test dimension in range (GridDim::maxThreadsPerBlock, GridDim::maxThreadsPerBlock * 2], ie. [257, 512]
TEST(TopKOperator, SmallArrayTopKSorted) {
std::vector<float> input_vals(400, 0.0f);
std::iota(input_vals.begin(), input_vals.end(), 0.0f);
std::vector<int64_t> input_dimensions = {400};
std::vector<float> expected_vals(400, 0.0f);
std::iota(expected_vals.begin(), expected_vals.end(), 0.0f);
std::reverse(expected_vals.begin(), expected_vals.end());
std::vector<int64_t> expected_indices(400, 0);
std::iota(expected_indices.begin(), expected_indices.end(), 0);
std::reverse(expected_indices.begin(), expected_indices.end());
std::vector<int64_t> expected_dimensions = {400};
RunTest(11, 400, input_vals, input_dimensions, expected_vals, expected_indices, expected_dimensions, false, -1, 1, 1);
}
TEST(TopKOperator, MediumArrayTopKSorted) {
std::vector<float> input_vals(1000, 0.0f);
std::iota(input_vals.begin(), input_vals.end(), 0.0f);
std::vector<int64_t> input_dimensions = {1000};
@ -490,8 +504,7 @@ TEST(TopKOperator, MediumArrayTopKSorted)
RunTest(11, 100, input_vals, input_dimensions, expected_vals, expected_indices, expected_dimensions, false, 0, 1, 1);
}
TEST(TopKOperator, BigArrayTopKSorted)
{
TEST(TopKOperator, BigArrayTopKSorted) {
std::vector<float> input_vals(10000, 0.0f);
std::iota(input_vals.begin(), input_vals.end(), 0.0f);
std::vector<int64_t> input_dimensions = {10000};
@ -505,8 +518,7 @@ TEST(TopKOperator, BigArrayTopKSorted)
RunTest(11, 1000, input_vals, input_dimensions, expected_vals, expected_indices, expected_dimensions, false, 0, 1, 1);
}
TEST(TopKOperator, BigArrayBigTopKSorted)
{
TEST(TopKOperator, BigArrayBigTopKSorted) {
std::vector<float> input_vals(10000, 0.0f);
std::iota(input_vals.begin(), input_vals.end(), 0.0f);
std::vector<int64_t> input_dimensions = {10000};