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 84.78261% with 7 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/onnx/parse_topk.cpp	57.14%	6 Missing ⚠️
src/program.cpp	66.67%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           develop    #4880      +/-   ##
===========================================
+ Coverage    92.86%   92.87%   +0.01%     
===========================================
  Files          585      585              
  Lines        30152    30213      +61     
===========================================
+ Hits         27998    28059      +61     
  Misses        2154     2154

Files with missing lines	Coverage Δ
src/include/migraphx/dyn_output.hpp	`100.00% <100.00%> (ø)`
src/include/migraphx/op/topk.hpp	`98.72% <100.00%> (+0.28%)`	⬆️
src/rewrite_topk.cpp	`100.00% <100.00%> (ø)`
src/program.cpp	`81.05% <66.67%> (-0.08%)`	⬇️
src/onnx/parse_topk.cpp	`78.57% <57.14%> (-21.43%)`	⬇️

... and 12 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.

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

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 27 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!

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));

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!

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