[AIROCMLIR-798] Add LDS usage estimate CAPI function

[justinrosner]

Motivation

Adds a new public CAPI, mlirMIGraphXLDSUsageFitsArch, that lets MIGraphX cheaply check (from problem sizes alone) whether a fused GEMM+GEMM / attention problem's estimated peak LDS (shared-memory) usage fits within a target arch's capacity.

Note: A version of this PR will also be opened up on rocmlirTriton so that both repos use the same CAPI version.

Technical Details

• New CAPI: bool mlirMIGraphXLDSUsageFitsArch(int64_t gemmO, const char *arch, MlirType elementType).
• New estimator rock::estimateGemmGemmLdsBytes(): computes the peak LDS footprint by evaluating both phases of the fused kernel (gemm0 A/B input tiles vs. the staged gemm0→gemm1 tile + gemm1's V tile) against the default attention perf config, accounting for double-buffered schedules and sub-byte element types.
• Extracted the data-type normalization logic out of ParamLookupTable::getDataTypeString() into a rock::getSupportedDataTypeString() , so the estimator can reject unsupported types gracefully
• Centralized the default attention perf-config string as rock::kDefaultAttnPerfConfig and consumed it from AffixTuningParameters.cpp, the new estimator, and the rocmlir-gen comment that previously duplicated the literal.

Note: Once this gets approval I will open up a corresponding rocmlirTriton, and MIGraphX PR so that the CAPI version is consistent across all three.

Test Plan

• MIGraphX CI
• PR CI

Test Result

• MIGraphX CI
• PR CI

Submission Checklist

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

[codecov]

Codecov Report

❌ Patch coverage is 77.46479% with 16 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
mlir/lib/Dialect/Rock/utility/loweringUtils.cpp	83.87%	4 Missing and 6 partials ⚠️
mlir/lib/Dialect/Rock/Tuning/ParamLookupTable.cpp	25.00%	5 Missing and 1 partial ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           develop    #2400      +/-   ##
===========================================
+ Coverage    82.57%   82.74%   +0.17%     
===========================================
  Files          120      120              
  Lines        42852    42872      +20     
  Branches      7110     7112       +2     
===========================================
+ Hits         35381    35472      +91     
+ Misses        4815     4794      -21     
+ Partials      2656     2606      -50     

Flag	Coverage Δ
gfx120x	82.67% <74.65%> (+0.15%)	⬆️
gfx950	82.41% <74.65%> (+0.06%)	⬆️
mfma	82.42% <77.46%> (?)

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

Files with missing lines	Coverage Δ
.../Dialect/Rock/Transforms/AffixTuningParameters.cpp	84.41% <100.00%> (-0.06%)	⬇️
mlir/tools/rocmlir-gen/rocmlir-gen.cpp	86.55% <ø> (ø)
mlir/lib/Dialect/Rock/Tuning/ParamLookupTable.cpp	87.85% <25.00%> (+2.37%)	⬆️
mlir/lib/Dialect/Rock/utility/loweringUtils.cpp	85.53% <83.87%> (+0.78%)	⬆️

... and 29 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[bdevorem]

I expected you to need all of the dims, m, n, k, and gemmO. Is the latter the only value that affects tile width?

[justinrosner]

m, n, and k don't really matter because we are going to choose tiling sizes based off the perf_config. However, for the second gemm, we calculate gemm1NPerBlock = llvm::PowerOf2Ceil(gemmO) (which matches what we do in rocMLIR and rocmlirTriton). So it's really a combination of the tile sizes from the perfConfig and gemmO that impacts the amount of LDS used.

[bdevorem]

if the tuned config is much different, will that have any kind of impact here?

[justinrosner]

I purposely went with a very small tile size here. If there is not enough LDS for this perfConfig with whatever gemmO size is passed in, then there is not going to be enough for any of the potentially more performant perfConfigs.

[dhernandez0]

I wonder if it makes sense to make the name a bit more generic, in case we want to extend it in the future. Something like mlirMIGraphXGemmGemmFusionFaster?

[justinrosner]

I was initially thinking about this the other way around (i.e., what if there are potentially more uses for a LDS check outside of GEG in MIGraphX in the future). I'm happy to make the name change to make it more GEG specific, but it just seems like we might have to make API updates regardless if both expanded use cases happen.

[bdevorem]

from my MIGraphX perspective, mlirMIGraphXLDSUsageFitsArch would be preferable for clarity

[Copilot]

Pull request overview

This PR introduces a new public MIGraphX C API to conservatively estimate whether a fused GEMM+GEMM/attention problem’s peak LDS usage fits a target GPU architecture, and centralizes the default attention perf-config string so it can be shared by multiple call sites.

Changes:

• Adds mlirMIGraphXLDSUsageFitsArch() C API and a corresponding CAPI test exercising fit/non-fit and invalid input cases.
• Implements rock::estimateGemmGemmLdsBytes() plus shared datatype normalization (rock::getSupportedDataTypeString()).
• Centralizes the default attention perf config as rock::kDefaultAttnPerfConfig and updates users to reference it.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
mlir/tools/rocmlir-gen/rocmlir-gen.cpp	Updates comment to reference centralized default attention perf config.
mlir/test/CAPI/mixr_full.c	Adds runtime test coverage for the new LDS-usage C API.
mlir/lib/Dialect/Rock/utility/loweringUtils.cpp	Adds datatype normalization helper and LDS-usage estimator.
mlir/lib/Dialect/Rock/Tuning/ParamLookupTable.cpp	Reuses shared datatype normalization; improves fatal error message with printed type.
mlir/lib/Dialect/Rock/Transforms/AffixTuningParameters.cpp	Switches default attention perf config to kDefaultAttnPerfConfig.
mlir/lib/CAPI/Dialect/MIGraphX.cpp	Adds mlirMIGraphXLDSUsageFitsArch() implementation.
mlir/include/mlir/Dialect/Rock/utility/loweringUtils.h	Declares the new helper + estimator functions.
mlir/include/mlir/Dialect/Rock/Tuning/GridwiseGemmGemmParams.h	Defines kDefaultAttnPerfConfig.
mlir/include/mlir-c/Dialect/MIGraphX.h	Bumps MIGraphX C API version and declares the new C API.

---

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

[umangyadav]

nit:
could there be cases of reduction fusion on the output where it utilizes more than peakBits estimated here ?

[justinrosner]

Yes, this is a valid concern. It could return true but then error out later on when rocMLIR actually compiles for certain reduction fusion cases. @bdevorem at the time of GEG fusion do you know if reductions are going to be added as well? If I were to add a bool hasOutputReduction parameter I think I could update the heuristic to also cover the case of reduction fusions

[bdevorem]

are you guys talking about reduces that get fused into a g+g fused submodule, or also reduction fusions that follow g+g fusions (i.e. small ones that do not get fused into the submodule)?

Either way, I should be able to check if those fusions happen down the pipe during g+g fusion

[justinrosner]

We are talking about reductions that get fused into a GEG kernel (i.e., in the same submodule)

[justinrosner]

So I was taking a deeper look into this and here is what I found. In the simple case where we have a single reduction that is the same dtype as the second GEMM, the reduction is not going to change the maxLDS answer already given by the function currently.

Things start to change a little bit when the reduction is in a wider element type or multiple reductions are fused and concurrently live. However, this liveness decision happens in ReuseLDS so the best that we could do is conservatively assume that all reductions are going to overlap? I don't love the idea of doing that, but @umangyadav do you have any opinions here?

[pfultz2]

I agree with @bdevorem here, I think we should have an initial gate using gemm sizes. We can handle the reduce fusion with this later.

[umangyadav]

How would MIGraphX find out whether Fusion can fit into LDS space or not later ?
Would it be using the same API that is being added here ? If so we need to modify the API to accept module instead of just operator. IMO we need to discuss design on this first.

[justinrosner]

@pfultz2 @bdevorem would it be acceptable from a MIGraphX point of view to have mlirMIGraphXLDSUsageFitsArch also take an optional mlir module? This way the function could serve two purposes. If used early on in the flow (without a mlir module passed in) it could still give you an estimate about LDS usage, but then the same function could be used later on in additional fusion checks when a module is present to get a definitive answer about whether or not things will fit into LDS.

[bdevorem]

Yes this works for me. We can still discuss in the interlock today if needed though

[justinrosner]

@bdevorem I update the mlirMIGraphXLDSUsageFitsArch function so that it can take an optional MlirModule. The function can now be used early in the flow when no module is yet present, and later on (to get a more definitive answer) when a module exists)

[rocmlir-pr-reviewer]

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

---

Scope

Adds a public CAPI mlirMIGraphXLDSUsageFitsArch (MIGraphX dialect API bumped to v6) that estimates peak LDS usage of a fused GEMM+GEMM/attention problem and checks it against an arch's shared-memory capacity. Introduces rock::estimateGemmGemmLdsBytes, extracts rock::getSupportedDataTypeString out of ParamLookupTable::getDataTypeString, and centralizes the default attention perf-config string as rock::kDefaultAttnPerfConfig. Touches the MIGraphX CAPI header/impl, Rock tuning/utility code, rocmlir-gen, and the mixr_full.c CAPI test.

Findings

• mlir/lib/CAPI/Dialect/MIGraphX.cpp:108 — the caller-provided arch is forwarded to estimateGemmGemmLdsBytes → lookupArchInfo, which hits llvm_unreachable for any unknown arch (AmdArchDb.cpp:414), so an invalid arch aborts the process instead of returning false as the header documents (Major: robustness/contract correctness; cf. checklist "Missing assert... use llvm_unreachable for impossible paths" — caller input is not an impossible path).
• mlir/include/mlir/Dialect/Rock/Tuning/GridwiseGemmGemmParams.h:26 — comment typo "available.." (Minor: comments as English prose).

Notes

• Copilot's earlier comment that getNumStages won't compile under C++17 (comparing std::optional<GemmLoadTileType> against an enum) is incorrect: std::optional's heterogeneous operator== is valid C++17, and the "C/C++ premerge checks" status is green, so no change is warranted there.
• The new CAPI is covered positively/negatively in mixr_full.c (valid arch, oversized problem, unsupported type, non-positive gemmO, null arch) but there is no test for a syntactically-valid-but-unknown arch string, which is exactly the path that currently crashes — worth adding once the validation above is in place.
• The estimator's reduction-fusion / tuned-config concerns raised inline by reviewers are acknowledged design trade-offs (conservative small-tile gate) and not blocking.

CI status

A Jenkins entry and ml-ci-internal.amd.com show as failing in the snapshot; several pipeline checks are still pending. These are not the auto-review pipeline's own check — please confirm the Jenkins failure is unrelated before merge.

[rocmlir-pr-reviewer]

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

---

Scope

Adds the public CAPI mlirMIGraphXLDSUsageFitsArch (and supporting rock::estimateGemmGemmLdsBytes / rock::getSupportedDataTypeString helpers) so MIGraphX can cheaply gate a fused GEMM+GEMM/attention problem on estimated peak LDS usage. Also centralizes the default attention perf-config string as rock::kDefaultAttnPerfConfig and bumps MLIR_MIGRAPHX_DIALECT_API_VERSION to 6.

Findings

No blocking issues found. Both previously-flagged items are resolved in this revision.

Notes

• The prior arch-abort concern (unknown arch reaching lookupArchInfo's llvm_unreachable) is now guarded at the CAPI boundary via RocmDeviceName::parse (MIGraphX.cpp:113-118), with a badarch negative case added to mixr_full.c. The remaining well-formed-but-unsupported case (e.g. a syntactically valid but unmodeled gfx*) is now explicitly documented in the header as a caller contract / UB. Note that rock::estimateGemmGemmLdsBytes itself (public in loweringUtils.h) still calls lookupArchInfo without validation, so any future non-CAPI caller must perform the same arch validation; worth a follow-up if the helper gains other callers.
• The bit-domain footprint computation with llvm::divideCeil(peakBits, 8) correctly avoids per-element rounding for sub-byte types. The getNumStages std::optional == enum comparison is valid in C++17 (the earlier Copilot note was a false positive, as the author noted).
• CAPI test coverage is solid: positive fit, oversized, unsupported dtype, non-positive gemmO, null arch, and malformed arch are all exercised.

CI status

The pipeline's own review check is in-progress (expected). detect shows as cancelled and the C/C++ and Python premerge checks were still in progress at snapshot time — none indicate a code-level failure attributable to this review.