Add dynamic shape support for TopK by klin2024 · Pull Request #4880 · ROCm/AMDMIGraphX

klin2024 · 2026-05-13T06:51:29Z

Motivation

The topk operator previously required a constant k and static input shapes, which will block the SSDMobileNetV2 model from running.

AMDMIGraphX/src/onnx/parse_topk.cpp

Line 47 in 93b8849

check_arg_empty(arg_k, "PARSE_TopK: k input must be constant");

This PR adds dynamic input support for the topk op.

Affected model: SSDMobileNetV2

Technical Details

.\bin\migraphx-driver.exe perf .\bin\topk_shape_derived_k.onnx can run success.

onnx file
topk_shape_derived_k.zip

Add a CHANGELOG.md entry for any option other than Not Applicable

- Added: New functionality.
- Changed: Changes to existing functionality.
- Removed: Functionality or support that has been removed. (Compared to a previous release)
- Optimized: Component performance that has been optimized or improved.
- Resolved Issues: Known issues from a previous version that have been resolved.
- Not Applicable: This PR is not to be included in the changelog.

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

codecov · 2026-05-13T09:48:01Z

Codecov Report

❌ Patch coverage is 97.87234% with 1 line in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/onnx/parse_topk.cpp	90.91%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           develop    #4880      +/-   ##
===========================================
- Coverage    92.86%   92.66%   -0.20%     
===========================================
  Files          585      588       +3     
  Lines        30152    30433     +281     
===========================================
+ Hits         27998    28199     +201     
- Misses        2154     2234      +80

Files with missing lines	Coverage Δ
src/include/migraphx/dyn_output.hpp	`100.00% <100.00%> (ø)`
src/include/migraphx/op/topk.hpp	`98.77% <100.00%> (+0.33%)`	⬆️
src/program.cpp	`81.61% <100.00%> (+0.48%)`	⬆️
src/rewrite_topk.cpp	`100.00% <100.00%> (ø)`
src/onnx/parse_topk.cpp	`95.65% <90.91%> (-4.35%)`	⬇️

... and 27 files with indirect coverage changes

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Copilot

Pull request overview

This PR extends MIGraphX’s TopK support to handle dynamic input shapes and cases where k is not a compile-time constant (by using a placeholder k derived from the input shape), and adjusts dynamic-output handling in evaluation/GPU lowering to better support tuple-shaped outputs like TopK.

Changes:

Update ONNX TopK parsing to allow non-constant k by deriving a placeholder k from the input shape (static lens or dynamic max lens).
Update op::topk to support dynamic shapes at shape-inference/eval time via dyn_output, including runtime clamping of k.
Adjust GPU/lowering + dyn-output plumbing for tuple/dynamic shapes, and add tests for dynamic TopK scenarios.

Reviewed changes

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

Show a summary per file

File	Description
test/verify/test_topk_dynamic.cpp	Adds a verify test exercising TopK with a dynamic input shape and placeholder `k`.
test/ref/topk.cpp	Adds reference tests for `k > n` with dynamic input and for `k == n` behavior.
src/targets/gpu/lowering.cpp	Switches dynamic-shape detection to `any_of_dynamic()` to account for tuple outputs.
src/targets/gpu/include/migraphx/gpu/hip.hpp	Minor whitespace-only adjustment.
src/targets/gpu/compile_ops.cpp	Alters dynamic code-object execution to always compute/reshape runtime output shape and adds a skip-on-empty heuristic.
src/rewrite_topk.cpp	Disables large-TopK rewrite when input shape is dynamic.
src/program.cpp	Makes tracing more robust when evaluating submodules from dynamic code-object ops.
src/onnx/parse_topk.cpp	Removes the “k must be constant” restriction; derives placeholder `k` from input shape when needed.
src/include/migraphx/op/topk.hpp	Adds dynamic-shape support and switches compute to `dyn_output`, clamping `k` at runtime.
src/include/migraphx/dyn_output.hpp	Uses `any_of_dynamic()` so tuple outputs with dynamic subshapes get runtime-computed shapes.

+                auto input_shape = args.at(0)->get_shape();
+                auto ndim        = input_shape.ndim();
+                auto norm_axis   = axis < 0 ? axis + static_cast<int64_t>(ndim) : axis;
+                if(input_shape.dynamic())
+                {


+                // k is not constant: use the input dimension along the topk axis
+                auto input_shape = args.at(0)->get_shape();
+                auto ndim        = input_shape.ndim();
+                auto norm_axis   = axis < 0 ? axis + static_cast<int64_t>(ndim) : axis;
+                if(input_shape.dynamic())
+                {
+                    k = input_shape.dyn_dims().at(norm_axis).get_interval().max;
+                }
+                else
+                {
+                    k = input_shape.lens().at(norm_axis);


+        auto out_shape = pre_op.compute_shape(to_shapes(static_args), module_args);
+        static_args[static_args.size() - 1] = output_arg.reshape(out_shape);
+        // Skip JIT compilation when dynamic shape resolves to 0 elements at runtime
+        if(args.front().get_shape().elements() == 0)
+            return static_args.back();


        auto topk_ret = info.add_instruction(
            make_op("topk", {{"k", k}, {"axis", axis}, {"largest", largest}}), args.at(0));



Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

wendyli-2024 · 2026-05-26T03:02:51Z

Hi @pfultz2
Do you have time to check the code change? Thanks

CharlieL7 · 2026-05-27T17:55:46Z

+            if(arg_k.empty())
+            {
+                // k is not constant: use the input dimension along the topk axis


Rather than a placeholder value of max_len make the k attribute in topk a std::optional<int64_t>. This would be more obvious what is meant.

CharlieL7 · 2026-05-27T17:59:09Z

        auto ins   = r.result;
        auto input = ins->inputs().front();
+        if(input->get_shape().dynamic())
+            return;


Use not_dynamic_shape in the matcher rather than checking here.

… represent unknown k instead of using a placeholder value

github-actions · 2026-06-01T06:15:53Z

-        lens[axis] = k;
+        if(inputs.at(0).dynamic())
+        {
+            auto dyn_dims     = inputs.at(0).dyn_dims();


[format.py] _{reported by reviewdog 🐶}

Suggested change

auto dyn_dims = inputs.at(0).dyn_dims();

auto dyn_dims = inputs.at(0).dyn_dims();

github-actions · 2026-06-01T06:15:53Z

        argument res_ind{vec_ss.back()};
        auto in_val       = args.front();
        auto relements    = in_val.get_shape().lens()[axis];
+        auto actual_k     = k.has_value() ? std::min(static_cast<std::size_t>(*k), relements) : relements;


[format.py] _{reported by reviewdog 🐶}

Suggested change

auto actual_k = k.has_value() ? std::min(static_cast<std::size_t>(*k), relements) : relements;

auto actual_k =

k.has_value() ? std::min(static_cast<std::size_t>(*k), relements) : relements;

github-actions · 2026-06-01T06:15:53Z

+                : info.add_instruction(
+                      make_op("topk", {{"axis", axis}, {"largest", largest}}), args.at(0));


[format.py] _{reported by reviewdog 🐶}

Suggested change

: info.add_instruction(

make_op("topk", {{"axis", axis}, {"largest", largest}}), args.at(0));

: info.add_instruction(make_op("topk", {{"axis", axis}, {"largest", largest}}),

args.at(0));

github-actions · 2026-06-01T06:15:53Z

+    auto matcher() const
+    {
+        return match::name("topk")(match::arg(0)(match::not_dynamic_shape()));
+    }


[format.py] _{reported by reviewdog 🐶}

Suggested change

auto matcher() const

{

return match::name("topk")(match::arg(0)(match::not_dynamic_shape()));

}

auto matcher() const { return match::name("topk")(match::arg(0)(match::not_dynamic_shape())); }

github-actions · 2026-06-01T06:15:53Z

 argument hip_topk::compute(context& ctx, const shape&, const std::vector<argument>& args) const
 {
    auto outputs = args.back().get_sub_objects();
+    auto actual_k = op.k.has_value() ? *op.k : static_cast<int64_t>(args[0].get_shape().lens()[op.axis]);


[format.py] _{reported by reviewdog 🐶}

Suggested change

auto actual_k = op.k.has_value() ? *op.k : static_cast<int64_t>(args[0].get_shape().lens()[op.axis]);

auto actual_k =

op.k.has_value() ? *op.k : static_cast<int64_t>(args[0].get_shape().lens()[op.axis]);

kazhang2 and others added 12 commits May 11, 2026 20:36

Support dynamic input for Topk op

02ec17c

fix debug version print and use internal allocation for dynamic output

b998457

add bound check

3527d11

Update src/include/migraphx/op/topk.hpp

171aec9

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Update src/include/migraphx/dyn_output.hpp

222a2dd

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Update src/program.cpp

33398d8

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Update src/replace_allocate.cpp

b067618

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Update src/targets/gpu/compile_ops.cpp

36ad89d

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Update test/verify/test_topk_dynamic.cpp

b55ae07

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Update test/verify/test_topk_dynamic.cpp

df26e31

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Update src/targets/gpu/compile_ops.cpp

89f31ba

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Update src/targets/gpu/include/migraphx/gpu/hip.hpp

2624cec

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

pfultz2 requested changes May 13, 2026

View reviewed changes

Comment thread src/targets/gpu/include/migraphx/gpu/hip.hpp Outdated

Comment thread src/replace_allocate.cpp Outdated

kazhang2 and others added 12 commits May 13, 2026 23:42

remove lifetime::global and s.any_of_dynamic check

bf9e4e4

fix license check

e4073bd

remove lifetime::global

eedd15f

format change

49b757b

Update test/ref/topk.cpp

3499f43

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Update src/targets/gpu/lowering.cpp

d5d0d3e

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Update test/ref/topk.cpp

68332cc

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Update test/ref/topk.cpp

0ee21ec

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Update test/ref/topk.cpp

3431d57

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Update test/ref/topk.cpp

c360580

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Update test/verify/test_topk_dynamic.cpp

e6090e3

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Update src/targets/gpu/include/migraphx/gpu/hip.hpp

688864a

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

klin2024 changed the title ~~Topk op support dynamic input fix~~ Add dynamic shape support for TopK May 14, 2026

klin2024 marked this pull request as ready for review May 14, 2026 15:45

klin2024 requested a review from causten as a code owner May 14, 2026 15:45

Copilot AI review requested due to automatic review settings May 14, 2026 15:45

Merge branch 'develop' into topk_op_support_dynamic_input_fix

13b06cb

Copilot AI reviewed May 14, 2026

View reviewed changes

causten requested a review from pfultz2 May 14, 2026 17:31

kazhang2 and others added 2 commits May 16, 2026 09:32

Potential fix for pull request finding

a048213

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Potential fix for pull request finding

0d08c50

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

CharlieL7 requested review from CharlieL7 and shivadbhavsar May 26, 2026 18:21

CharlieL7 requested changes May 27, 2026

View reviewed changes

CharlieL7 reviewed May 27, 2026

View reviewed changes

Change topk op k from int64_t to std::optional<int64_t> to explicitly…

040e186

… represent unknown k instead of using a placeholder value

github-actions Bot reviewed Jun 1, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add dynamic shape support for TopK#4880

Add dynamic shape support for TopK#4880
klin2024 wants to merge 28 commits into
developfrom
topk_op_support_dynamic_input_fix

klin2024 commented May 13, 2026 •

edited

Loading

Uh oh!

codecov Bot commented May 13, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

wendyli-2024 commented May 26, 2026

Uh oh!

CharlieL7 May 27, 2026

Uh oh!

CharlieL7 May 27, 2026

Uh oh!

github-actions Bot Jun 1, 2026

Uh oh!

github-actions Bot Jun 1, 2026

Uh oh!

github-actions Bot Jun 1, 2026

Uh oh!

github-actions Bot Jun 1, 2026

Uh oh!

github-actions Bot Jun 1, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

		auto topk_ret = info.add_instruction(
		make_op("topk", {{"k", k}, {"axis", axis}, {"largest", largest}}), args.at(0));

	auto dyn_dims = inputs.at(0).dyn_dims();
	auto dyn_dims = inputs.at(0).dyn_dims();

	auto actual_k = k.has_value() ? std::min(static_cast<std::size_t>(*k), relements) : relements;
	auto actual_k =
	k.has_value() ? std::min(static_cast<std::size_t>(*k), relements) : relements;

		: info.add_instruction(
		make_op("topk", {{"axis", axis}, {"largest", largest}}), args.at(0));

	auto actual_k = op.k.has_value() ? *op.k : static_cast<int64_t>(args[0].get_shape().lens()[op.axis]);
	auto actual_k =
	op.k.has_value() ? *op.k : static_cast<int64_t>(args[0].get_shape().lens()[op.axis]);

Conversation

klin2024 commented May 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Technical Details

Uh oh!

codecov Bot commented May 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

wendyli-2024 commented May 26, 2026

Uh oh!

CharlieL7 May 27, 2026

Choose a reason for hiding this comment

Uh oh!

CharlieL7 May 27, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jun 1, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jun 1, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jun 1, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jun 1, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jun 1, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

klin2024 commented May 13, 2026 •

edited

Loading

codecov Bot commented May 13, 2026 •

edited

Loading