Add squeeze, gather, and scatter_none ops to MIGraphX dialect #2176

justinrosner · 2025-12-17T20:53:32Z

Motivation

This PR adds three new ops to the MIGraphX dialect (squeeze, gather, and scatter_none). These new ops are required as part of the paged attention work.

This implements https://github.com/ROCm/rocMLIR-internal/issues/2200

Technical Details

This PR can be broken down into the following changes:

MIGraphX.td:
- Adds migraphx.squeeze (https://onnx.ai/onnx/operators/onnx__Squeeze.html): Remove dimensions of size 1
- Adds migraphx.gather (https://onnx.ai/onnx/operators/onnx__Gather.html): Gather slices from data along an axis
- Adds migraphx.scatter_none (https://onnx.ai/onnx/operators/onnx__ScatterElements.html): Updates into data at specified indices
MIGraphX.cpp: Add verifiers for each of the new ops
MIGraphXToTosa: Lowers each of the new ops to their set of TOSA ops

Test Plan

Add new LIT tests for each of the new ops that was added

Test Result

PR CI

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

justinrosner · 2025-12-17T22:13:07Z

Note for reviewers, this PR is operating under the assumption that we will have the upstream changes in soon so that we can have this change: llvm/llvm-project#167894 (adding i64 type support for indices in scatter/gather). To be conservative, I can keep the conversion from i64 -> i32 and just add a TODO and file a ticket to make that simple change once the upstream PR makes its way in. Let me know what you guys want.

Copilot

Pull request overview

This PR adds three new operations to the MIGraphX dialect as part of the paged attention work: squeeze, gather, and scatter_none. These ops are based on ONNX specifications and include complete implementations with verifiers, TOSA lowering patterns, and comprehensive test coverage.

Implements ONNX-compatible squeeze (removing size-1 dimensions), gather (collecting slices along an axis), and scatter_none (updating elements at specified indices)
Adds verifiers for axis bounds checking, shape validation, and index constraints
Provides TOSA lowering implementations that reshape tensors to match TOSA's 3D gather/scatter semantics

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 8 comments.

Show a summary per file

File	Description
mlir/include/mlir/Dialect/MIGraphX/IR/MIGraphX.td	Adds TableGen definitions for squeeze, gather, and scatter_none ops with documentation and assembly formats
mlir/lib/Dialect/MIGraphX/IR/MIGraphX.cpp	Implements verify() methods for the three new ops including axis bounds checking and shape validation
mlir/lib/Conversion/MIGraphXToTosa/MIGraphXToTosa.cpp	Adds converter patterns to lower each new op to TOSA operations via reshape/gather/scatter transformations
mlir/test/Dialect/MIGraphX/invalid.mlir	Adds negative test cases for axis out of bounds, rank mismatches, and invalid indices
mlir/test/Conversion/MIGraphXToTosa/migraphx-to-tosa.mlir	Adds positive conversion tests for squeeze, gather, and scatter_none with various axis configurations

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

mlir/lib/Dialect/MIGraphX/IR/MIGraphX.cpp

mlir/lib/Conversion/MIGraphXToTosa/MIGraphXToTosa.cpp

mlir/include/mlir/Dialect/MIGraphX/IR/MIGraphX.td

mlir/lib/Conversion/MIGraphXToTosa/MIGraphXToTosa.cpp

mlir/include/mlir/Dialect/MIGraphX/IR/MIGraphX.td

mlir/lib/Dialect/MIGraphX/IR/MIGraphX.cpp

pfultz2 · 2026-01-06T21:01:34Z

mlir/include/mlir/Dialect/MIGraphX/IR/MIGraphX.td

  }];
 }

+def MIGraphX_SqueezeOp


You dont need to add squeeze and unsqueeze as we already lower this to reshape when going to mlir.

The IR that I was basing this off of (which came from MIGraphX) was using the unlowered squeeze/unsqueeze. I think it's a fairly trivial lowering, so it wouldn't hurt if rocMLIR could also do something like this.

pfultz2 · 2026-01-06T21:02:53Z

Outside of paged attention, will this enable us to fuse gather and scatter into gemms and convolutions?

justinrosner · 2026-01-06T21:05:46Z

Outside of paged attention, will this enable us to fuse gather and scatter into gemms and convolutions?

This PR in of itself won't get us to fusing gather/scatter operations. We would still need an additional backend lowering PR for this.

dhernandez0 · 2026-01-07T09:27:56Z

mlir/include/mlir/Dialect/MIGraphX/IR/MIGraphX.td

+  let hasVerifier = 1;
+}
+
+def MIGraphX_ScatterNoneOp


why is it called ScatterNone?

MIGraphX was passing us a scatter_none op which directly correlates to the ScatterElements ONNX op with reduction set to none: https://onnx.ai/onnx/operators/onnx__ScatterElements.html

dhernandez0 · 2026-01-07T09:36:02Z

mlir/lib/Conversion/MIGraphXToTosa/MIGraphXToTosa.cpp

+    }
+  }
+
+  // Build the new shape by excluding the squeezed axes


should we check that axesToSqueeze are all dimension=1?

The verifier already does this validation.

Or is that not enough? Since if I recall correctly the verifier isn't turned on by default?

dhernandez0 · 2026-01-07T09:38:18Z

mlir/lib/Conversion/MIGraphXToTosa/MIGraphXToTosa.cpp

+  auto outputType = cast<RankedTensorType>(outputTy);
+  Type elemType = dataType.getElementType();
+
+  // Lowering strategy for migraphx.gather -> tosa.gather:


why not skip tosa instead of this workaround?

I was thinking about needing a lowering path through TOSA for the CPU path in the future (there currently isn't a lowering path that goes past Tosa). I've already opened up a ticket to address this: https://github.com/ROCm/rocMLIR-internal/issues/2205. It looks like IREE has already encountered this before and they have a custom lowering that maybe we could port/use.

dhernandez0 · 2026-01-07T09:39:29Z

mlir/lib/Conversion/MIGraphXToTosa/MIGraphXToTosa.cpp

+  // Output: [N, W, C] where each [n, w, :] = reshaped_data[n, indices[n, w], :]
+  SmallVector<int64_t> gatherOutputShape = {N, W, C};
+  auto gatherOutputType = RankedTensorType::get(gatherOutputShape, elemType);
+  Value gatherResult = tosa::GatherOp::create(rewriter, loc, gatherOutputType,


I understand we are missing tosa.gather -> rock? are we going to support this for paged attention for now?

The tosa.gather -> rock was going to come in a future PR that handled the TosaToRock changes. However, in recent conversations regarding paged attention, it seems like we are moving away from the scatter/gather implementation.

do we need this PR then?

@pfultz2 Are there any plans for MIGraphX to require rocMLIR supporting gather/scatter outside of paged attention (think I might have heard something in one of the meetings yesterday)? If not, then I think we can close this.

Fusing gather for gemms and convolutions will be useful as well, especially for resize.

dhernandez0 · 2026-01-07T09:39:44Z

mlir/lib/Conversion/MIGraphXToTosa/MIGraphXToTosa.cpp

+  if (axis < 0)
+    axis += dataRank;
+
+  // TOSA scatter requires that indices be constant across the "C" dimension


same here, why not skip tosa?

See comment above about CPU lowering.

mlir/lib/Dialect/MIGraphX/IR/MIGraphX.cpp

justinrosner added 4 commits December 17, 2025 18:53

Add MIGraphX dialect ops

abfc318

MIGraphXToTosa changes

fbcd035

Clang-format

4a488a0

Add LIT tests

af8d9f5

justinrosner added 4 commits December 18, 2025 15:24

Small bug fixes

9d960e9

Clang-format

0eca6ed

Fix up comment

32d8833

Remove unnecessary negative index check

7a78ffa

justinrosner marked this pull request as ready for review December 18, 2025 15:42

justinrosner requested a review from causten as a code owner December 18, 2025 15:42

justinrosner requested review from Copilot, dhernandez0, pabloantoniom and umangyadav December 18, 2025 15:42

Copilot started reviewing on behalf of justinrosner December 18, 2025 15:47 View session

Copilot AI reviewed Dec 18, 2025

View reviewed changes

Int fixes

fec3854

justinrosner changed the title ~~[WIP] Add squeeze, gather, and scatter_none ops to MIGraphX dialect~~ Add squeeze, gather, and scatter_none ops to MIGraphX dialect Jan 2, 2026

pfultz2 reviewed Jan 6, 2026

View reviewed changes

dhernandez0 reviewed Jan 7, 2026

View reviewed changes

justinrosner added 2 commits January 8, 2026 18:52

Attend to review comments

597ac36

Clang-format

a4462ca

justinrosner force-pushed the 2200-migraphx-paged-attention-ops branch from 37703a4 to a4462ca Compare January 8, 2026 18:55

Merge branch 'develop' into 2200-migraphx-paged-attention-ops

9098300

justinrosner requested a review from dhernandez0 January 12, 2026 21:51

Add squeeze, gather, and scatter_none ops to MIGraphX dialect #2176

Are you sure you want to change the base?

Add squeeze, gather, and scatter_none ops to MIGraphX dialect #2176

Uh oh!

Conversation

justinrosner commented Dec 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Uh oh!

justinrosner commented Dec 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

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!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pfultz2 commented Jan 6, 2026

Uh oh!

justinrosner commented Jan 6, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

justinrosner commented Dec 17, 2025 •

edited

Loading

justinrosner commented Dec 17, 2025 •

edited

Loading