saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-27 20:02:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Support int32 type for TopK CPU op (#7089)

Browse source
This commit is contained in:
Hariharan Seshadri 2021-03-31 18:08:21 -07:00 committed by GitHub
parent b370ddbf5e
commit 0e0dd50e39
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 19 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
onnxruntime/core/providers/cpu/cpu_execution_provider.cc
View file

@ -418,6 +418,7 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, Un
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, float, TopK);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, double, TopK);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, int64_t, TopK);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, int32_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);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, int64_t_string_int64_t, OneHot);
@ -1272,6 +1273,8 @@ Status RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) {
TopK)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, int64_t,
TopK)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, int32_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,

11
onnxruntime/core/providers/cpu/math/top_k.cc
View file

@ -491,6 +491,11 @@ TopK<11, double>::TopK(const OpKernelInfo& op_kernel_info) : OpKernel(op_kernel_
TopkOpset11ConstructorCommon(op_kernel_info, axis_, largest_, sorted_);
}
template <>
TopK<11, int32_t>::TopK(const OpKernelInfo& op_kernel_info) : OpKernel(op_kernel_info) {
TopkOpset11ConstructorCommon(op_kernel_info, axis_, largest_, sorted_);
}
template <>
TopK<11, int64_t>::TopK(const OpKernelInfo& op_kernel_info) : OpKernel(op_kernel_info) {
TopkOpset11ConstructorCommon(op_kernel_info, axis_, largest_, sorted_);
@ -507,6 +512,11 @@ Status TopK<11, double>::Compute(OpKernelContext* p_op_kernel_context) const {
return ComputeImplOpset1011<double>(p_op_kernel_context, axis_, largest_, sorted_);
}
template <>
Status TopK<11, int32_t>::Compute(OpKernelContext* p_op_kernel_context) const {
return ComputeImplOpset1011<int32_t>(p_op_kernel_context, axis_, largest_, sorted_);
}
template <>
Status TopK<11, int64_t>::Compute(OpKernelContext* p_op_kernel_context) const {
return ComputeImplOpset1011<int64_t>(p_op_kernel_context, axis_, largest_, sorted_);
@ -539,5 +549,6 @@ REGISTER_TOPK_VERSIONED_TYPED_KERNEL(10, 10, double);
REGISTER_TOPK_TYPED_KERNEL(11, float);
REGISTER_TOPK_TYPED_KERNEL(11, double);
REGISTER_TOPK_TYPED_KERNEL(11, int64_t);
REGISTER_TOPK_TYPED_KERNEL(11, int32_t);
} // namespace onnxruntime

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

@ -584,7 +584,7 @@ TEST(TopKOperator, Top3ExplicitAxisSmallestElements) {
template <typename T>
static void top_1_explicit_axis_MultiD_input_smallest(int opset_version, int64_t sorted = 1) {
std::vector<T> input_vals = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f};
std::vector<T> input_vals = {1, 2, 3, 4, 5, 6, 7, 8};
std::vector<int64_t> input_dimensions = {2, 2, 2};
std::vector<T> expected_vals = {1, 2, 5, 6};
std::vector<int64_t> expected_indices = {0, 0, 0, 0};
@ -598,6 +598,10 @@ TEST(TopKOperator, Top1ExplicitAxisMultiDInputSmallestElements) {
top_1_explicit_axis_MultiD_input_smallest<float>(11, 0); //unsorted
top_1_explicit_axis_MultiD_input_smallest<double>(11);
top_1_explicit_axis_MultiD_input_smallest<double>(11, 0); //unsorted
top_1_explicit_axis_MultiD_input_smallest<int32_t>(11);
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
}
// test path where SelectTopK is used (select using std::nth_element)