onnxruntime
7fc826a4 - Add INT8, INT16, and UINT8 type support for CPU TopK operator (#27860)

Commit
6 days ago
Add INT8, INT16, and UINT8 type support for CPU TopK operator (#27860) The ONNX specification (opset 11+) lists INT8, INT16, and UINT8 as valid input types for the TopK operator, but ONNX Runtime's CPU execution provider only registered float, double, int32, and int64 kernels. This causes a NOT_IMPLEMENTED error when running models that produce TopK nodes with these smaller integer types (e.g., PP-DocLayoutV2 exported via torch.onnx.export(dynamo=True)). This commit adds kernel registrations and template specializations for int8_t, int16_t, and uint8_t in opset 11-23 and opset 24, along with unit tests covering largest, smallest, negative values, explicit axis, and opset 24 scenarios. ### Description Add kernel registrations and template specializations for `int8_t`, `int16_t`, and `uint8_t` types in the CPU TopK operator (opset 11-23 and opset 24). Fixes #27859 **Changed files:** - `onnxruntime/core/providers/cpu/math/top_k.cc` — template specializations + kernel registrations - `onnxruntime/core/providers/cpu/cpu_execution_provider.cc` — forward declarations + BuildKernelCreateInfo entries - `onnxruntime/test/providers/cpu/math/topk_op_test.cc` — 8 new test cases (largest, smallest, negative values, explicit axis, opset 24) ### Motivation and Context The [ONNX specification (opset 11+)](https://github.com/onnx/onnx/blob/main/docs/Operators.md#TopK) lists `INT8`, `INT16`, and `UINT8` as valid input types for the TopK operator. However, ONNX Runtime's CPU execution provider only had kernels registered for `float`, `double`, `int32`, and `int64`, causing a `NOT_IMPLEMENTED` error: ``` ONNXRuntimeError: Could not find an implementation for TopK(11) node with name '...' ``` This issue is encountered when running models like PP-DocLayoutV2 exported via `torch.onnx.export(dynamo=True)`, which produces a `Cast(bool → INT8) → TopK` pattern. The existing CPU TopK implementation is fully template-based, so no algorithmic changes were needed — only kernel registration and template instantiation for the missing types. All 64 TopK tests pass (including 8 new tests for the added types).
Author
Parents
Loading