[royzhang-amd]

Motivation

Graphs with an ArgMax reduction followed by a squeeze can fail during reshape fusion with SQUEEZE: axis dimension should be equal to 1. For example, ArgMax operations can reduce axis 1 of inputs shaped {1, 2304, 1} and {1, 2304, 17}, followed by a squeeze of the reduced axis.

Reshape fusion can split the reduction dimension into common dimensions such as 48 x 48. ArgMin and ArgMax return indices into the original axis, so reducing only one component changes the index semantics and can leave a non-singleton dimension for the downstream squeeze. Using the first combined reduction axis is unambiguous only when the fused reduction has one axis and that axis remains whole after remapping.

This change limits arg-reduction reshape fusion to one aggregate reduction axis that maps to exactly one common axis. Multi-axis or split-axis cases skip the rewrite.

Technical Details

• reduce_reshape::supports: for a fused module containing ArgMin or ArgMax, require axes.size() == 1 and exactly one common axis in its mapping. This permits supported transpose remapping and skips multi-axis or split-axis rewrites before graph mutation.

• Use the existing source and destination support checks in rewrite_reshapes before inserting instructions or replacing the graph. The rebase incorporates the newer upstream reshape-fusion implementation.

• Keep the existing axes.front() remapping in reduce_reshape::insert. The preflight guard guarantees one mapped reduction axis; the remaining attributes, including select_last_index, are preserved.

• test/fuse_reduce.cpp: compare complete expected programs for split-axis/squeeze rejection, supported source and destination transposes, and destination split rejection. ArgMin and ArgMax regressions also assert that a multi-axis fused module followed by a nontrivial transpose stays unchanged. Check preservation of select_last_index=true.

• test/verify/test_arg_reduce_split_axis_squeeze.cpp: verify ArgMin and ArgMax, first/last index selection, and tied/varied runtime inputs against the reference target. Use make_op for construction. Bound runtime addends with int8 inputs converted to int64, retaining exact integer output checks while avoiding precision loss in the CPU floating fallback and overflow in verification arithmetic.

• Scope: the restriction applies to fused modules containing arg reductions. Multi-axis and split-axis cases skip reshape fusion without throwing; the original graph remains available for compilation.

Validation on a Windows Radeon 880M (gfx1150) board:

• Initial validation before this rebase reproduced the original compilation failure and the regression failure with the pre-fix production code.

• Before this rebase, the fixed FP32 model passed all seven deterministic input cases against ONNX Runtime CPU with both default and serial scheduling. CPU fallback was disabled; profiles showed only MIGraphX execution. Output shape and type checks and exact repeated-output checks passed. Maximum absolute error was 5.14e-6, within atol=1e-4, rtol=1e-3 for every output element. These checks use synthetic inputs, not a labeled pose-accuracy dataset.

• The original fix passed 208 targeted tests: 60 test_fuse_reduce, 52 test_fuse_pointwise, and 96 test_shape_transform_descriptor cases.

• Initial cleanup validation: all 11 selected arg* unit cases and all 16 new GPU/reference verification cases passed against the validated fixed runtime.

• Before this rebase, additional small ONNX models covered source/destination axis splits, supported remapping, transpose, and singleton dimensions. All 24 IR checks and 48 numerical comparisons passed. Last-index cases used direct MIGraphX GPU/reference verification because the installed ONNX Runtime MIGraphX EP rejects select_last_index=1 during capability checks; the pre-fix runtime has the same restriction.

• Previous revision 09a89d1e0 passed 219 targeted unit tests (71 test_fuse_reduce, 52 test_fuse_pointwise, 96 test_shape_transform_descriptor).

• On revision 09a89d1e0, all eight verification variants matched exactly when replaying CPU unsupported-type lowering on the reference target; the original unbounded inputs reproduced eight mismatches in the same check. That verification graph is unchanged by the single-axis restriction. This replay does not replace a full CPU backend run.

• Current single-axis-only revision: all 72 test_fuse_reduce cases passed. Mutation checks confirmed failures when the aggregate single-axis condition is removed, split axes are accepted, or every arg-reduction rewrite is rejected. Restoring the candidate passed all 72 cases again.

• Formatting with CI's pinned clang-format 22.1.5 and git diff --check passed. Full CPU/GPU CI and static analysis for the updated commit remain pending.

Changelog Category

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.

Follow the LLVM AI Tool Use Policy for contributions using AI.

[Copilot]

🟡 Changes recommended

Committed regressions still need to cover successful per-operation remapping and destination-side rejection.

2 open issues

• Add fusion tests for axis remapping and select_last_index · New
• Add destination-side reshape guard regression test · New

What changed in this PR

Preserves ArgMin/ArgMax index semantics during MIGraphX reshape fusion, avoiding invalid split-axis rewrites.

Changes:

• Checks source and destination support before modifying the graph.
• Remaps supported arg-reduction axes individually while preserving attributes.
• Updates expected graphs and adds split-axis/squeeze regression coverage.

File	Description
test/​verify/​test_arg_reduce_split_axis_squeeze.cpp	Adds GPU/reference cases for ties and index-selection modes.
test/​fuse_reduce.cpp	Updates expected graphs and adds squeeze regressions.
src/​include/​migraphx/​rewrite_reshapes.hpp	Checks support before rewriting either operation.
src/​fuse_reduce.cpp	Rejects unsupported axis mappings and remaps each arg reduction.

_{Review effort: Balanced}