Add TopK layer support for pnnx ONNX export path

[vlordier]

Summary

• add native TopK layer implementation in ncnn core
• lower ONNX TopK nodes in pnnx ncnn pass to the new TopK layer
• register layer/build integration and add test_topk coverage

Why

Issue #6377 reports ONNX export pipelines (for example YOLOv10 via ultralytics) failing due to unsupported TopK in ncnn conversion/runtime path.

Validation

• configured and built test_topk locally
• ran ./build/tests/test_topk successfully

Fixes #6377

[tencent-adm]

All committers have signed the CLA.

[Copilot]

Pull request overview

This pull request adds TopK layer support for the pnnx ONNX export path in ncnn, addressing issue #6377 which reports failures in ONNX export pipelines for models like YOLOv10 through ultralytics.

Changes:

• Implements native TopK layer in ncnn core with support for different axes, largest/smallest selection, and sorted/unsorted output
• Adds pnnx pass to convert ONNX TopK nodes to ncnn TopK layer with proper batch axis handling
• Adds comprehensive test coverage for 1D, 2D, 3D, and 4D tensors with various axis configurations

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
src/layer/topk.h	Header file defining TopK layer class with parameters for axis, largest, sorted, and k
src/layer/topk.cpp	Implementation of TopK layer supporting 1-4D tensors with efficient sorting using std::partial_sort and std::nth_element
src/CMakeLists.txt	Registers TopK layer in build system (correctly ordered alphabetically)
tools/pnnx/src/pass_ncnn/TopK.cpp	PNNX pass to convert ONNX TopK operations to ncnn format, handling batch axis removal and parameter mapping
tools/pnnx/src/CMakeLists.txt	Adds TopK pass to pnnx build (correctly ordered alphabetically)
tests/test_topk.cpp	Test cases covering 1D through 4D tensors with various axis, k, and largest parameters
tests/CMakeLists.txt	Registers TopK test (minor ordering issue - should come after Tile)

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[github-actions]

Please enable github action in YOUR FORKED REPO to make code-format workflow work

[mlknz]

any updates on this?

[vlordier]

@mlknz seems it's passing the tests and CI, give it a try and let me know

[Copilot]

Pull request overview

Copilot reviewed 22 out of 22 changed files in this pull request and generated 12 comments.

Comments suppressed due to low confidence (1)

src/layer/cast.cpp:108

• Cast advertises type_from=0 as “auto” (see cast.h), but Cast::forward() never resolves type_from==0 to an actual source type. With the new pnnx Tensor.to -> Cast lowering emitting type_from=0, this will produce uninitialized outputs. Add type_from==0 inference (e.g., from bottom_blob.elemsize / known flags) or explicitly treat 0 as float32 when no better signal exists, and return an error for unsupported conversions.

int Cast::forward(const Mat& bottom_blob, Mat& top_blob, const Option& opt) const
{
    if (type_from == type_to)
    {
        top_blob = bottom_blob;
        return 0;
    }

    int w = bottom_blob.w;
    int h = bottom_blob.h;
    int d = bottom_blob.d;
    int channels = bottom_blob.c;
    int dims = bottom_blob.dims;
    size_t elemsize = bottom_blob.elemsize;
    int elempack = bottom_blob.elempack;

    size_t out_elemsize = elemsize;
    if (type_to == 1)
    {
        // float32
        out_elemsize = 4 * elempack;
    }
    else if (type_to == 2)
    {
        // float16
        out_elemsize = 2 * elempack;
    }
    else if (type_to == 3)
    {
        // int8
        out_elemsize = elempack;
    }
    else if (type_to == 4)
    {
        // bfloat16
        out_elemsize = 2 * elempack;
    }
    else if (type_to == 5)
    {
        // int64
        out_elemsize = 8 * elempack;
    }
    else if (type_to == 6)
    {
        // int32
        out_elemsize = 4 * elempack;
    }

    if (dims == 1)
    {
        top_blob.create(w, out_elemsize, elempack, opt.blob_allocator);
    }
    else if (dims == 2)
    {
        top_blob.create(w, h, out_elemsize, elempack, opt.blob_allocator);
    }
    else if (dims == 3)
    {
        top_blob.create(w, h, channels, out_elemsize, elempack, opt.blob_allocator);
    }
    else if (dims == 4)
    {
        top_blob.create(w, h, d, channels, out_elemsize, elempack, opt.blob_allocator);
    }
    if (top_blob.empty())
        return -100;

    int size = w * h * d * elempack;

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Missing semicolon after return 0 causes a compile error, blocking the new test_topk target from building. Add the semicolon so the chained || test_topk(...) expressions compile.

Suggested change

	return 0
	|| test_topk(a, 0, 1, 1, 1)
	|| test_topk(a, 0, 5, 1, 1)
	|| test_topk(a, 0, 1, 0, 0)
	|| test_topk(a, -1, 7, 0, 1)
	|| test_topk(a, 0, 4, 1, 0)
	|| test_topk(a, 0, 9, 1, 1);
	}
	static int test_topk_1()
	{
	ncnn::Mat a = RandomMat(12, 17);
	return 0
	|| test_topk(a, 0, 1, 1, 1)
	|| test_topk(a, 0, 5, 1, 1)
	|| test_topk(a, 1, 3, 1, 1)
	|| test_topk(a, -1, 8, 0, 1)
	|| test_topk(a, 1, 6, 0, 0)
	|| test_topk(a, -2, 7, 1, 1);
	}
	static int test_topk_2()
	{
	ncnn::Mat a = RandomMat(8, 9, 11);
	return 0
	|| test_topk(a, 0, 3, 1, 1)
	|| test_topk(a, 1, 4, 1, 1)
	|| test_topk(a, 2, 2, 0, 1)
	|| test_topk(a, 2, 5, 1, 0)
	|| test_topk(a, -1, 6, 1, 1)
	|| test_topk(a, -2, 5, 0, 1)
	|| test_topk(a, -3, 7, 1, 1);
	return 0 ||
	test_topk(a, 0, 1, 1, 1) ||
	test_topk(a, 0, 5, 1, 1) ||
	test_topk(a, 0, 1, 0, 0) ||
	test_topk(a, -1, 7, 0, 1) ||
	test_topk(a, 0, 4, 1, 0) ||
	test_topk(a, 0, 9, 1, 1);
	}
	static int test_topk_1()
	{
	ncnn::Mat a = RandomMat(12, 17);
	return 0 ||
	test_topk(a, 0, 1, 1, 1) ||
	test_topk(a, 0, 5, 1, 1) ||
	test_topk(a, 1, 3, 1, 1) ||
	test_topk(a, -1, 8, 0, 1) ||
	test_topk(a, 1, 6, 0, 0) ||
	test_topk(a, -2, 7, 1, 1);
	}
	static int test_topk_2()
	{
	ncnn::Mat a = RandomMat(8, 9, 11);
	return 0 ||
	test_topk(a, 0, 3, 1, 1) ||
	test_topk(a, 1, 4, 1, 1) ||
	test_topk(a, 2, 2, 0, 1) ||
	test_topk(a, 2, 5, 1, 0) ||
	test_topk(a, -1, 6, 1, 1) ||
	test_topk(a, -2, 5, 0, 1) ||
	test_topk(a, -3, 7, 1, 1);

[Copilot]

The PNNXIR header 12 7 does not match the actual number of operator lines in this pattern (7). This mismatch will prevent the graph rewriter from matching aten::topk nodes. Update the header counts to reflect the pattern contents (likely 7 7).

[Copilot]

When instantiating TopK modules, this loop prints it.first as a Python keyword (e.g., 0=..., 1=...). ncnn-style param keys are numeric strings, which are not valid Python identifiers and will generate invalid Python code. Map these param ids to named arguments (axis, largest, sorted, k) or emit positional arguments in a fixed order.

[Copilot]

Gather::forward() reinterprets the index tensor as const int* (idx = (const int*)index_blob), but ncnn (and this PR’s TopK/ArgMax) commonly represent indices as float values stored in a float buffer. Reinterpreting float bits as int will produce huge indices and out-of-bounds reads. Read indices as float and cast to int (or branch based on index_blob.elemsize / the Cast output type).

Suggested change

	const int* idx = (const int*)index_blob;
	const float* idx = index_blob;

[Copilot]

This implementation only decomposes output coordinates / computes flat_in for dims up to 3. For 4D inputs (and/or 4D index tensors), coord_out/flat_in are left partially unset and flat_in stays 0, so most outputs will incorrectly read the first element. Add explicit handling for dims==4 (both coordinate decomposition and flat index computation, including cstep/d strides).

[Copilot]

Tensor.to is lowered to Cast with type_from=0 (“auto”), but Cast::forward() currently has no logic to infer the source type when type_from==0, so none of the conversion branches will run and the output buffer will be left uninitialized. Either implement type_from==0 auto-detection in Cast or set an explicit type_from here based on the actual input dtype tracked by pnnx.

[Copilot]

The ncnn TopK layer expects k via param id 3 (or via a second input blob), but this rewrite pass never writes k into op->params. As a result, converted graphs will use the default k=1 regardless of the original torch.topk(..., k=...). Capture and write k (and consider supporting the ONNX-style second input for dynamic K if needed).

[Copilot]

The generated Python TopK module takes forward(self, x, k), but the pnnx->ncnn lowering for torch.topk in this PR models k as a captured parameter (not an input tensor). If the resulting TopK operator only has one input, the emitted self.<name>(v_input) call will raise at runtime due to the missing k argument. Align the Python stub with how k is represented (store k as a module attribute and accept only x, or ensure the operator keeps a second k input).

Suggested change

	fprintf(pyfp, " def __init__(self, axis=1, largest=1, sorted=1):\n");
	fprintf(pyfp, " super(TopK, self).__init__()\n");
	fprintf(pyfp, " self.axis = axis\n");
	fprintf(pyfp, " self.largest = largest\n");
	fprintf(pyfp, " self.sorted = sorted\n");
	fprintf(pyfp, " def forward(self, x, k):\n");
	fprintf(pyfp, " def __init__(self, axis=1, largest=1, sorted=1, k=None):\n");
	fprintf(pyfp, " super(TopK, self).__init__()\n");
	fprintf(pyfp, " self.axis = axis\n");
	fprintf(pyfp, " self.largest = largest\n");
	fprintf(pyfp, " self.sorted = sorted\n");
	fprintf(pyfp, " self.k = k\n");
	fprintf(pyfp, " def forward(self, x, k=None):\n");
	fprintf(pyfp, " k = self.k if k is None else k\n");
	fprintf(pyfp, " if k is None:\n");
	fprintf(pyfp, " raise TypeError('TopK.forward() missing required argument: k')\n");

[Copilot]

top_blob.create(out_shape.w, out_shape.h, out_shape.c, ...) always creates a 3D output, even when index_blob is 1D/2D/4D. That changes the output rank and breaks the coordinate decomposition logic below (which depends on top_blob.dims). Allocate top_blob using the correct create() overload based on out_shape.dims so the output shape truly matches the index tensor.

Suggested change

	top_blob.create(out_shape.w, out_shape.h, out_shape.c, input_blob.elemsize, input_blob.elempack, opt.blob_allocator);
	if (out_shape.dims == 1)
	top_blob.create(out_shape.w, input_blob.elemsize, input_blob.elempack, opt.blob_allocator);
	else if (out_shape.dims == 2)
	top_blob.create(out_shape.w, out_shape.h, input_blob.elemsize, input_blob.elempack, opt.blob_allocator);
	else if (out_shape.dims == 3)
	top_blob.create(out_shape.w, out_shape.h, out_shape.c, input_blob.elemsize, input_blob.elempack, opt.blob_allocator);
	else if (out_shape.dims == 4)
	top_blob.create(out_shape.w, out_shape.h, out_shape.d, out_shape.c, input_blob.elemsize, input_blob.elempack, opt.blob_allocator);
	else
	return -1;

[Copilot]

This lowering writes axis directly from dim, but other ncnn rewrite passes drop/renumber the batch axis using __batch_index (see e.g. torch_flip/TopK). Without the same adjustment here, Gather will target the wrong dimension after batch removal. Apply the same batch-axis handling (reject gather along batch, and decrement axes > batch_index).

Suggested change

	if (op->params.find("__batch_index") != op->params.end())
	{
	const Parameter& batch_index_p = op->params.at("__batch_index");
	if (batch_index_p.type == 2)
	{
	const int batch_index = batch_index_p.i;
	// ncnn lowers tensor ranks with the batch axis removed.
	// gather along the removed batch axis cannot be represented here.
	if (axis == batch_index)
	return;
	if (axis > batch_index)
	axis -= 1;
	}
	}

[vlordier]

Superseded by #6669 which consolidates all operators (TopK, Gather, GatherElements, Expand, Tile, Mod) with all review issues addressed.