[NFC] Add elementwise trait on MIGraphX dialect ops

[umangyadav]

Cherry picking 6b1119a

Will be useful when doing migraphx.attention op

Also rearranges migraphx/ops.mlir and migraphx/invalid.mlir tests so that all tests for a operation appear all together.

Also adds type contraints on the migraphx.clip op.

[Copilot]

Pull request overview

This PR cherry-picks changes that mark MIGraphX dialect elementwise operations with MLIR’s Elementwise trait (and adds/strengthens some shape-equality verification for unary/convert ops), while reorganizing and expanding MIGraphX dialect tests to keep per-op coverage grouped together.

Changes:

• Add Elementwise trait to MIGraphX elementwise binary/unary ops (and to where, clip, convert), plus add AllShapesMatch constraints for unary ops and convert.
• Expand mlir/test/Dialect/MIGraphX/ops.mlir with grouped, per-op smoke tests for many elementwise ops.
• Reorganize mlir/test/Dialect/MIGraphX/invalid.mlir and add new negative tests for shaped-type stride/shape mismatches and elementwise op verification.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File	Description
mlir/test/Dialect/MIGraphX/ops.mlir	Reorders and adds per-op CHECK blocks for elementwise ops (plus existing dot/quant_dot coverage).
mlir/test/Dialect/MIGraphX/invalid.mlir	Regroups invalid tests by op/type and adds new verifier diagnostics coverage.
mlir/include/mlir/Dialect/MIGraphX/IR/MIGraphX.td	Adds Elementwise trait broadly and introduces AllShapesMatch for unary ops and convert.

Comments suppressed due to low confidence (1)

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

• The migraphx.convert op description says it "currently only supports float to float conversions", but the tree contains tests and conversion pipelines that use migraphx.convert for integer conversions as well (e.g., MIGraphXToLinalg tests include f32<->i32 cases). Please either update the description to match current supported behavior, or add/enforce a verifier restriction if non-float conversions are truly unsupported.

def MIGraphX_ConvertOp
    : MIGraphX_Op<"convert", [Elementwise, AllShapesMatch<["inA", "output"]>]>,
      Arguments<(ins AnyMIXRShaped:$inA)>,
      Results<(outs AnyMIXRShaped:$output)> {
  let summary = "Elementwise type conversion";
  let description = [{
    Type conversion. Due to impedance mismatches between MIGraphX and Tosa,
    currently only supports float to float conversions
  }];

---

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

[dhernandez0]

from #2316 (comment)

why limit preSoftmaxBody to elementwise only? I think they might have reshapes inside the presoftmax currently. This might be too strict.

why do we need this trait?

[umangyadav]

from #2316 (comment)
why do we need this trait?

I've not run into cases where preSoftmaxBody has reshapes in it. I can update logic later if needed. This PR is only concerned about marking migraphx ops as elementwise.

In general it is good idea to mark them as elementwise so that we can do pattern matching more generically.

[dhernandez0]

I've not run into cases where preSoftmaxBody has reshapes in it. I can update logic later if needed. This PR is only concerned about marking migraphx ops as elementwise.

we do have cases like that, for example mixr-attention-gqa2.mlir, see how there's at least a reshape between the first and second gemm.

In general it is good idea to mark them as elementwise so that we can do pattern matching more generically.

please, can you provide more context for this PR? it's difficult to see the usefulness with the code here.

[umangyadav]

we do have cases like that, for example mixr-attention-gqa2.mlir, see how there's at least a reshape between the first and second gemm.

Yes, i see them. Thanks.

please, can you provide more context for this PR? it's difficult to see the usefulness with the code here.

I want to verify that preSoftmaxBody has elementwise ops and reshape ops. If i don't have elementwiseTrait then i'll have to keep a long list in Op verifier for preSoftmaxBody.

e.g. i am doing this right now, see how it uses ElementwiseTrait. If not i would have to keep a list of elementwise ops there.

  Region &body = getPreSoftmaxBody();
  bool hasPreSoftmaxInputs = !getPreSoftmaxElemWiseInputs().empty();
  bool hasNonTerminatorOps = false;
  for (Block &block : body) {
    for (Operation &op : block) {
      if (op.hasTrait<OpTrait::IsTerminator>())
        continue;
      hasNonTerminatorOps = true;
      Dialect *dialect = op.getDialect();
      bool isMIGraphX = dialect && dialect->getNamespace() == "migraphx";
      if (!isMIGraphX || !op.hasTrait<OpTrait::Elementwise>())
        return op.emitOpError(
                   "preSoftmaxBody must only contain elementwise migraphx ops, "
                   "but found '")
               << op.getName() << "'";
    }
  }

[justinrosner]

Changes look good to me. Are there any plans on doing something similar for reshape ops? Or is that list short enough where it doesn't matter?

[umangyadav]

Or is that list short enough where it doesn't matter?

It's most probably going to be short list. If not i can define a trait later