[AIROCMLIR-1093] Support arbitrary element types for expand_strides

[justinrosner]

Motivation

Certain MIGraphX E2E models have quantized convolutions with non-contiguous i32 output layouts. The operation previously restricted operands to floating-point and i8 types, causing verification to reject valid i32 outputs.

Technical Details

This PR updates the expand_strides in the following ways:

• Allow rock.expand_strides to accept any tensor or memref element type.
• Use declarative traits to require matching operand ranks and element types.

Test Plan

• Nightly CI

Test Result

• Nightly CI

Submission Checklist

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

[Copilot]

Pull request overview

This PR broadens rock.expand_strides type support to unblock MIGraphX E2E models that produce non-contiguous i32 outputs (e.g., quantized convolutions). It does so by relaxing the operand element-type constraints while moving rank/element-type consistency checks to declarative ODS traits, and adds regression coverage.

Changes:

• Update rock.expand_strides to accept any tensor/memref element type, while requiring operand rank match + operand element-type match via declarative traits.
• Simplify ExpandStridesOp::verify() by removing redundant rank/element-type checks now covered by traits, keeping the shape-extent check.
• Add new Rock dialect tests (tensor + memref i32) and an E2E repro for non-contiguous i32 quant-conv output layouts.

Reviewed changes

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

Show a summary per file

File	Description
mlir/include/mlir/Dialect/Rock/IR/RockOps.td	Loosens expand_strides operand type constraints to any element type and adds declarative verification traits for rank + element-type matching.
mlir/lib/Dialect/Rock/IR/RockDialect.cpp	Removes now-redundant rank/element-type checks from ExpandStridesOp::verify(), leaving dimension-size validation.
mlir/test/Dialect/Rock/ops.mlir	Adds positive tests for rock.expand_strides on i32 tensors and memrefs.
mlir/test/Dialect/Rock/invalid.mlir	Updates expected diagnostics to match trait-driven verifier messages for rank/element-type mismatch cases.
mlir/test/fusion/pr-e2e/mixr-quant-conv-non-contiguous-strides.mlir	Adds an E2E regression test exercising the motivating i32 non-contiguous output-layout scenario.

---

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

[justinrosner]

This does not require a port to rocmlirTriton. See the comments here in the initial rocmlirTriton implementation for further details: https://github.com/ROCm/rocmlirTriton/pull/114

[rocmlir-pr-reviewer]

Verdict: APPROVE -- submitted as COMMENT (automated reviews are advisory)  ·  Findings: 0 (0 Critical, 0 Major, 0 Minor)

---

Scope

Relaxes rock.expand_strides to accept any tensor/memref element type (motivated by MIGraphX quantized convolutions with non-contiguous i32 outputs). Replaces the TensorOrMemRefOf<[F32,F16,BF16,I8]> operand constraint with AnyTensorOrMemRef, and moves the rank-match and element-type-match checks from the hand-written verify() into the declarative AllRanksMatch/SameOperandsElementType traits. Touches RockOps.td, RockDialect.cpp, two Lit tests, and adds one E2E fusion test.

Findings

No blocking issues found.

Notes

• The refactor is idiomatic MLIR: using declarative traits over hand-rolled checks satisfies the DRY guidance in the checklist. The removal of the manual rank check in RockDialect.cpp:880 is safe because trait verifiers run before the op's verify(), so the remaining llvm::zip_equal over shapes is guaranteed equal ranks first.
• The invalid.mlir diagnostics were correctly updated to the trait-generated wording, and ops.mlir adds both tensor and memref i32 positive cases. The op retains hasVerifier = 1 with a meaningful verify() body (output-dim >= input-dim), so the checklist's op-verifier requirement still holds.
• SameOperandsElementType constrains only the two operands, not the optional result; this matches the pre-existing behavior (the old verify() never checked the result element type either), so no regression is introduced.
• Reviewer @umangyadav already flagged moving mlir/test/fusion/pr-e2e/mixr-quant-conv-non-contiguous-strides.mlir to the nightly suite; worth resolving that thread before merge. A minor trailing blank line was added at the end of ops.mlir — harmless but easy to drop.

CI status

All functional checks pass (C/C++ premerge, Python format/lint, Python perf). The auto-review pipeline's own review check is expected to be in-progress and is not a CI failure.

[codecov]

Codecov Report

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

Files with missing lines	Patch %	Lines
mlir/lib/Dialect/Rock/IR/RockDialect.cpp	80.00%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           develop    #2434      +/-   ##
===========================================
+ Coverage    82.57%   83.51%   +0.95%     
===========================================
  Files          120      120              
  Lines        42852    42934      +82     
  Branches      7110     7138      +28     
===========================================
+ Hits         35381    35856     +475     
+ Misses        4815     4501     -314     
+ Partials      2656     2577      -79     

Flag	Coverage Δ
gfx120x	83.39% <80.00%> (+0.86%)	⬆️
gfx950	83.29% <80.00%> (+0.95%)	⬆️
mfma	83.36% <80.00%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
mlir/lib/Dialect/Rock/IR/RockDialect.cpp	77.60% <80.00%> (+8.95%)	⬆️

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