Skip to content

Add INT8, INT16, and UINT8 type support for CPU TopK operator #27860

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tianleiwu merged 4 commits into from
Apr 8, 2026
Merged

Add INT8, INT16, and UINT8 type support for CPU TopK operator #27860

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
06e6210
Add INT8, INT16, and UINT8 type support for CPU TopK operator
elwhyjay Mar 26, 2026
4168072
Fix clang-format: remove extra alignment space in TopK test
elwhyjay Mar 27, 2026
af5948a
Fix comment to clarify opset range for TopK specializations (11-23)
elwhyjay Apr 1, 2026
3e65d87
Update docs/OperatorKernels.md for CPU TopK int8/int16/uint8 support
elwhyjay Apr 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions docs/OperatorKernels.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -503,8 +503,8 @@ Do not modify directly.*
|||[1, 9]|**T** = tensor(float)|
|Tile|*in* input:**T**<br> *in* repeats:**T1**<br> *out* output:**T**<br><br>or<br><br>*in* input:**T**<br> *in* tiles:**T**<br> *in* axis:**T**<br> *out* output:**T**|13+|**T** = tensor(bool), tensor(double), tensor(float), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)<br/> **T1** = tensor(int64)|
|||[6, 12]|**T** = tensor(bool), tensor(double), tensor(float), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)<br/> **T1** = tensor(int64)|
|TopK|*in* X:**T**<br> *in* K:**tensor(int64)**<br> *out* Values:**T**<br> *out* Indices:**I**<br><br>or<br><br>*in* X:**T**<br> *out* Values:**T**<br> *out* Indices:**I**|24+|**I** = tensor(int64)<br/> **T** = tensor(double), tensor(float), tensor(int32), tensor(int64)|
|||[11, 23]|**I** = tensor(int64)<br/> **T** = tensor(double), tensor(float), tensor(int32), tensor(int64)|
|TopK|*in* X:**T**<br> *in* K:**tensor(int64)**<br> *out* Values:**T**<br> *out* Indices:**I**<br><br>or<br><br>*in* X:**T**<br> *out* Values:**T**<br> *out* Indices:**I**|24+|**I** = tensor(int64)<br/> **T** = tensor(double), tensor(float), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint8)|
|||[11, 23]|**I** = tensor(int64)<br/> **T** = tensor(double), tensor(float), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint8)|
|||10|**I** = tensor(int64)<br/> **T** = tensor(double), tensor(float)|
|||[1, 9]|**I** = tensor(int64)<br/> **T** = tensor(double), tensor(float)|
|Transpose|*in* data:**T**<br> *out* transposed:**T**|25+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(int16), tensor(int2), tensor(int32), tensor(int4), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint2), tensor(uint32), tensor(uint4), tensor(uint64), tensor(uint8)|
Expand Down
15 changes: 15 additions & 0 deletions onnxruntime/core/providers/cpu/cpu_execution_provider.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -587,6 +587,9 @@ class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOn
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 23, double, TopK);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 23, int64_t, TopK);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 23, int32_t, TopK);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 23, int8_t, TopK);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 23, int16_t, TopK);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 23, uint8_t, TopK);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, int64_t_int64_t_int64_t, OneHot);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, float_int64_t_int64_t, OneHot);
#if !defined(DISABLE_STRING_TYPE)
Expand Down Expand Up @@ -1455,6 +1458,9 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain,
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 24, double, TopK);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 24, int64_t, TopK);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 24, int32_t, TopK);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 24, int8_t, TopK);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 24, int16_t, TopK);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 24, uint8_t, TopK);

// Opset 25
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 25, Cast);
Expand Down Expand Up @@ -2281,6 +2287,12 @@ Status RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) {
int64_t, TopK)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 23,
int32_t, TopK)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 23,
int8_t, TopK)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 23,
int16_t, TopK)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 23,
uint8_t, TopK)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11,
int64_t_int64_t_int64_t, OneHot)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11,
Expand Down Expand Up @@ -3598,6 +3610,9 @@ Status RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 24, double, TopK)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 24, int64_t, TopK)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 24, int32_t, TopK)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 24, int8_t, TopK)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 24, int16_t, TopK)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 24, uint8_t, TopK)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 24, SplitToSequence)>,

// opset 25
Expand Down
14 changes: 13 additions & 1 deletion onnxruntime/core/providers/cpu/math/top_k.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -542,17 +542,23 @@ static void TopkOpset11ConstructorCommon(const OpKernelInfo& op_kernel_info,
return ComputeImplOpset1011<TYPE>(p_op_kernel_context, axis_, largest_, sorted_); \
}

// Generate specializations for opset 11 (used by versioned kernel 11-23)
// Generate specializations for opset 11-23 (used by versioned kernel 11-23)
TOPK_MODERN_OPSET_SPECIALIZATIONS(23, float);
TOPK_MODERN_OPSET_SPECIALIZATIONS(23, double);
TOPK_MODERN_OPSET_SPECIALIZATIONS(23, int32_t);
TOPK_MODERN_OPSET_SPECIALIZATIONS(23, int64_t);
TOPK_MODERN_OPSET_SPECIALIZATIONS(23, int8_t);
TOPK_MODERN_OPSET_SPECIALIZATIONS(23, int16_t);
TOPK_MODERN_OPSET_SPECIALIZATIONS(23, uint8_t);
Comment thread
tianleiwu marked this conversation as resolved.

// Generate specializations for opset 24 (used by current kernel 24+)
TOPK_MODERN_OPSET_SPECIALIZATIONS(24, float);
TOPK_MODERN_OPSET_SPECIALIZATIONS(24, double);
TOPK_MODERN_OPSET_SPECIALIZATIONS(24, int32_t);
TOPK_MODERN_OPSET_SPECIALIZATIONS(24, int64_t);
TOPK_MODERN_OPSET_SPECIALIZATIONS(24, int8_t);
TOPK_MODERN_OPSET_SPECIALIZATIONS(24, int16_t);
TOPK_MODERN_OPSET_SPECIALIZATIONS(24, uint8_t);

// Register necessary kernels
// spec https://github.com/onnx/onnx/blob/main/docs/Operators.md#TopK
Expand Down Expand Up @@ -582,10 +588,16 @@ REGISTER_TOPK_VERSIONED_TYPED_KERNEL(11, 23, float);
REGISTER_TOPK_VERSIONED_TYPED_KERNEL(11, 23, double);
REGISTER_TOPK_VERSIONED_TYPED_KERNEL(11, 23, int64_t);
REGISTER_TOPK_VERSIONED_TYPED_KERNEL(11, 23, int32_t);
REGISTER_TOPK_VERSIONED_TYPED_KERNEL(11, 23, int8_t);
REGISTER_TOPK_VERSIONED_TYPED_KERNEL(11, 23, int16_t);
REGISTER_TOPK_VERSIONED_TYPED_KERNEL(11, 23, uint8_t);

REGISTER_TOPK_TYPED_KERNEL(24, float);
REGISTER_TOPK_TYPED_KERNEL(24, double);
REGISTER_TOPK_TYPED_KERNEL(24, int64_t);
REGISTER_TOPK_TYPED_KERNEL(24, int32_t);
REGISTER_TOPK_TYPED_KERNEL(24, int8_t);
REGISTER_TOPK_TYPED_KERNEL(24, int16_t);
REGISTER_TOPK_TYPED_KERNEL(24, uint8_t);

} // namespace onnxruntime
81 changes: 81 additions & 0 deletions onnxruntime/test/providers/cpu/math/topk_op_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -621,6 +621,12 @@ TEST(TopKOperator, Top1ExplicitAxisMultiDInputSmallestElements) {
top_1_explicit_axis_MultiD_input_smallest<int32_t>(11, 0); // unsorted
top_1_explicit_axis_MultiD_input_smallest<int64_t>(11);
top_1_explicit_axis_MultiD_input_smallest<int64_t>(11, 0); // unsorted
top_1_explicit_axis_MultiD_input_smallest<int8_t>(11);
top_1_explicit_axis_MultiD_input_smallest<int8_t>(11, 0); // unsorted
top_1_explicit_axis_MultiD_input_smallest<int16_t>(11);
top_1_explicit_axis_MultiD_input_smallest<int16_t>(11, 0); // unsorted
top_1_explicit_axis_MultiD_input_smallest<uint8_t>(11);
top_1_explicit_axis_MultiD_input_smallest<uint8_t>(11, 0); // unsorted
}

// test path where SelectTopK is used (select using std::nth_element)
Expand Down Expand Up @@ -909,5 +915,80 @@ TEST(TopKOperator, SelectTopKThreaded) {
TestThreaded<double>(k, n, batch_size);
}

// Tests for INT8, INT16, UINT8 types (opset 11+)
TEST(TopKOperator, TopK_Int8) {
std::vector<int8_t> input_vals = {10, 30, 20, 40, 10, 30, 40, 20};
std::vector<int64_t> input_dimensions = {2, 4};
std::vector<int8_t> expected_vals = {40, 30, 40, 30};
std::vector<int64_t> expected_indices = {3, 1, 2, 1};
std::vector<int64_t> expected_dimensions = {2, 2};
RunTest(11, 2, input_vals, input_dimensions, expected_vals, expected_indices, expected_dimensions, false);
}

TEST(TopKOperator, TopK_Int8_Negative) {
std::vector<int8_t> input_vals = {-10, -30, -20, -40, -10, -30, -40, -20};
std::vector<int64_t> input_dimensions = {2, 4};
std::vector<int8_t> expected_vals = {-10, -20, -10, -20};
std::vector<int64_t> expected_indices = {0, 2, 0, 3};
std::vector<int64_t> expected_dimensions = {2, 2};
RunTest(11, 2, input_vals, input_dimensions, expected_vals, expected_indices, expected_dimensions, false);
}

TEST(TopKOperator, TopK_Int8_Smallest) {
std::vector<int8_t> input_vals = {10, 30, 20, 40, 10, 30, 40, 20};
std::vector<int64_t> input_dimensions = {2, 4};
std::vector<int8_t> expected_vals = {10, 20, 10, 20};
std::vector<int64_t> expected_indices = {0, 2, 0, 3};
std::vector<int64_t> expected_dimensions = {2, 2};
RunTest(11, 2, input_vals, input_dimensions, expected_vals, expected_indices, expected_dimensions, false, -1, 0);
}

TEST(TopKOperator, TopK_Int16) {
std::vector<int16_t> input_vals = {100, 300, 200, 400, 100, 300, 400, 200};
std::vector<int64_t> input_dimensions = {2, 4};
std::vector<int16_t> expected_vals = {400, 300, 400, 300};
std::vector<int64_t> expected_indices = {3, 1, 2, 1};
std::vector<int64_t> expected_dimensions = {2, 2};
RunTest(11, 2, input_vals, input_dimensions, expected_vals, expected_indices, expected_dimensions, false);
}

TEST(TopKOperator, TopK_Uint8) {
std::vector<uint8_t> input_vals = {10, 30, 20, 40, 10, 30, 40, 20};
std::vector<int64_t> input_dimensions = {2, 4};
std::vector<uint8_t> expected_vals = {40, 30, 40, 30};
std::vector<int64_t> expected_indices = {3, 1, 2, 1};
std::vector<int64_t> expected_dimensions = {2, 2};
RunTest(11, 2, input_vals, input_dimensions, expected_vals, expected_indices, expected_dimensions, false);
}

TEST(TopKOperator, TopK_Int8_ExplicitAxis) {
std::vector<int8_t> input_vals = {1, 2, 3, 4, 5, 6, 7, 8};
std::vector<int64_t> input_dimensions = {2, 2, 2};
std::vector<int8_t> expected_vals = {3, 4, 7, 8};
std::vector<int64_t> expected_indices = {1, 1, 1, 1};
std::vector<int64_t> expected_dimensions = {2, 1, 2};
int64_t axis = 1;
RunTest(11, 1, input_vals, input_dimensions, expected_vals, expected_indices, expected_dimensions, false, axis);
}

// Opset 24 tests for new types
TEST(TopKOperator, TopK_Int8_Opset24) {
std::vector<int8_t> input_vals = {10, 30, 20, 40, 10, 30, 40, 20};
std::vector<int64_t> input_dimensions = {2, 4};
std::vector<int8_t> expected_vals = {40, 30, 40, 30};
std::vector<int64_t> expected_indices = {3, 1, 2, 1};
std::vector<int64_t> expected_dimensions = {2, 2};
RunTest(24, 2, input_vals, input_dimensions, expected_vals, expected_indices, expected_dimensions, false);
}

TEST(TopKOperator, TopK_Uint8_Opset24) {
std::vector<uint8_t> input_vals = {10, 30, 20, 40, 10, 30, 40, 20};
std::vector<int64_t> input_dimensions = {2, 4};
std::vector<uint8_t> expected_vals = {40, 30, 40, 30};
std::vector<int64_t> expected_indices = {3, 1, 2, 1};
std::vector<int64_t> expected_dimensions = {2, 2};
RunTest(24, 2, input_vals, input_dimensions, expected_vals, expected_indices, expected_dimensions, false);
}

} // namespace test
} // namespace onnxruntime
Loading