diff --git a/onnxruntime/core/providers/cuda/math/topk_impl.cu b/onnxruntime/core/providers/cuda/math/topk_impl.cu index fca3a5b06e..3ca5e1d982 100644 --- a/onnxruntime/core/providers/cuda/math/topk_impl.cu +++ b/onnxruntime/core/providers/cuda/math/topk_impl.cu @@ -22,6 +22,26 @@ struct KV { int64_t val; }; +template +struct NumericLimits { + static T Lowest() { + return std::numeric_limits::lowest(); + } + static T Max() { + return std::numeric_limits::max(); + } +}; + +template <> +struct NumericLimits { + static half Lowest() { + return -65504.0; + } + static half Max() { + return 65504.0; + } +}; + #define BT GridDim::maxThreadsPerBlock #define ALIGN(N) static_cast(pow(2, ceil(log2(static_cast(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 } } + template __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<<)>>>(input_x_ptr, output_v_ptr, output_i, elem_nums, size, axis, K, aligned_K, largest, sorted, dimension, aligned_dimension, std::numeric_limits::lowest(), std::numeric_limits::max()); + BitonicTopK<<)>>>(input_x_ptr, output_v_ptr, output_i, elem_nums, size, axis, K, aligned_K, largest, sorted, dimension, aligned_dimension, NumericLimits::Lowest(), NumericLimits::Max()); } else if (K <= BT*16 || 0 == sorted) { auto XPT = static_cast(ceil(static_cast(dimension) / GridDim::maxThreadsPerBlock)); if (BT*2 >= K || 0 == sorted) { - RadixTopK<<>>(input_x_ptr, output_v_ptr, output_i, elem_nums, size, axis, K, largest, sorted, dimension, XPT, std::numeric_limits::lowest(), std::numeric_limits::max()); + RadixTopK<<>>(input_x_ptr, output_v_ptr, output_i, elem_nums, size, axis, K, largest, sorted, dimension, XPT, NumericLimits::Lowest(), NumericLimits::Max()); } else if (BT*4>=K) { - RadixTopK<<>>(input_x_ptr, output_v_ptr, output_i, elem_nums, size, axis, K, largest, sorted, dimension, XPT, std::numeric_limits::lowest(), std::numeric_limits::max()); + RadixTopK<<>>(input_x_ptr, output_v_ptr, output_i, elem_nums, size, axis, K, largest, sorted, dimension, XPT, NumericLimits::Lowest(), NumericLimits::Max()); } else if (BT*8>=K) { - RadixTopK<<>>(input_x_ptr, output_v_ptr, output_i, elem_nums, size, axis, K, largest, sorted, dimension, XPT, std::numeric_limits::lowest(), std::numeric_limits::max()); + RadixTopK<<>>(input_x_ptr, output_v_ptr, output_i, elem_nums, size, axis, K, largest, sorted, dimension, XPT, NumericLimits::Lowest(), NumericLimits::Max()); } else { - RadixTopK<<>>(input_x_ptr, output_v_ptr, output_i, elem_nums, size, axis, K, largest, sorted, dimension, XPT, std::numeric_limits::lowest(), std::numeric_limits::max()); + RadixTopK<<>>(input_x_ptr, output_v_ptr, output_i, elem_nums, size, axis, K, largest, sorted, dimension, XPT, NumericLimits::Lowest(), NumericLimits::Max()); } } else { auto input_key_buffer = kernel->GetScratchBuffer(dimension); diff --git a/onnxruntime/test/providers/cpu/math/topk_op_test.cc b/onnxruntime/test/providers/cpu/math/topk_op_test.cc index 68f1eac63b..ef69c92c4a 100644 --- a/onnxruntime/test/providers/cpu/math/topk_op_test.cc +++ b/onnxruntime/test/providers/cpu/math/topk_op_test.cc @@ -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 input_vals_f = {10.0f, -8.0f, -7.0f, -4.0f, -5.0f, -6.0f}; + std::vector expected_vals_f = {10.0f, -4.0f, -5.0f, -6.0f}; + std::vector input_vals(6); + std::vector expected_vals(4); + ConvertFloatToMLFloat16(input_vals_f.data(), input_vals.data(), 6); + ConvertFloatToMLFloat16(expected_vals_f.data(), expected_vals.data(), 4); + std::vector input_dimensions = {6}; + std::vector expected_indices = {0, 3, 4, 5}; + std::vector 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 input_vals(400, 0.0f);