[AIROCMLIR-1067] Extend problemKey with supportsSplitK

[justinrosner]

Motivation

Expose whether a tuning problem supports Split-K so tuning infrastructure can avoid selecting Split-K configurations for incompatible fused operations. This fixes a problem where MIGraphX was caching a perfConfig for a problemKey, but that same problemKey matched two different kernels. One of which supported splitK > 1, and the other which did not.

rocmlirTriton port of this change: https://github.com/ROCm/rocmlirTriton/pull/430

Technical Details

• Append -supportsSplitK true|false to generated tuning keys.
• Determine support using testFusionLegalitySplitK.

Test Plan

• PR CI

Test Result

• PR CI

Submission Checklist

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

[Copilot]

Pull request overview

Extends rocMLIR tuning problem keys with explicit -supportsSplitK true|false metadata so downstream tuning/caching (e.g., MIGraphX) can distinguish otherwise-identical problems that differ in Split-K compatibility.

Changes:

• Appends -supportsSplitK {true|false} to serialized tuning keys and teaches perfRunner to parse/round-trip this metadata without passing it to rocmlir-gen / external tools.
• Computes Split-K support during tuning-key emission using rock::testFusionLegalitySplitK on the correct parent function (not the module’s first function).
• Updates Python and MLIR tests to reflect canonicalization behavior and to assert Split-K support tagging in emitted tuning keys.

Reviewed changes

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

Show a summary per file

File	Description
mlir/utils/performance/perfRunner.py	Adds tuning-key metadata extraction + round-trip support; strips metadata before invoking external tools where it’s not a valid option.
mlir/utils/performance/tests/test_tuningRunner.py	Updates canonicalization expectations and adds a round-trip test for supportsSplitK metadata.
mlir/utils/performance/tests/test_perfRunner.py	Updates tuning DB parsing expectations and adds a test ensuring Split-K support is distinguished in keys.
mlir/lib/Dialect/Rock/Tuning/RockTuningImpl.cpp	Appends -supportsSplitK to tuning keys based on testFusionLegalitySplitK evaluated on the serialized op’s parent function.
mlir/test/rocmlir-gen/gemm-misc-options.mlir	Updates emitted GEMM tuning-key expectation to include -supportsSplitK true.
mlir/test/fusion/fusability-dot-mul.mlir	Adds a tuning-key check asserting -supportsSplitK true for a Split-K-compatible fusion.
mlir/test/fusion/fusability-dot-add.mlir	Adds a tuning-key check asserting -supportsSplitK true for a Split-K-compatible fusion.
mlir/test/fusion/fusability-conv-add-max.mlir	Adds a tuning-key check asserting -supportsSplitK false for a Split-K-incompatible fusion.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[rocmlir-pr-reviewer]

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

---

Scope

Appends a -supportsSplitK true|false field to the tuning problem key so tuning infrastructure can distinguish two kernels that previously collapsed onto one problemKey. C++ side (RockTuningImpl.cpp) wraps the two getTuningProblemStr(ModuleOp) dispatch points in a lambda that also runs rock::testFusionLegalitySplitK on the enclosing func::FuncOp; Python side (perfRunner.py) adds PerfConfiguration.tuning_key_metadata(), threads a supports_split_k flag through all five config classes' from_command_line/to_command_line, and adds extract_tuning_key_metadata() so the new token is stripped before it can reach MIOpenDriver argv. Lit tests updated in mlir/test/fusion/ and mlir/test/rocmlir-gen/; pytest updated for the new canonical form.

Findings

• Critical — mlir/utils/performance/perfRunner.py:592: the Python default is False with no way to ever produce true from a config-file test vector, while the C++ emitter produces -supportsSplitK true for the same plain GEMM (see the updated GEMM check in mlir/test/rocmlir-gen/gemm-misc-options.mlir). The two key producers are explicitly required to stay byte-identical (# Please keep this in sync with mlir::rock::getTuningProblemStr() at perfRunner.py:759 and :1249), and today they do for GEMM. After this PR every row tuningRunner.py writes ends in -supportsSplitK false, so the runtime lookup never hits.
• Major — mlir/lib/Dialect/Rock/Tuning/RockTuningImpl.cpp:1753: unconditional key-format change with no fallback in tuningTableLookup, invalidating every pre-existing tuning-DB entry.
• Minor — RockTuningImpl.cpp:1754: testFusionLegalitySplitK builds a full BufferDependencyAnalysis; this now runs on every tuning-key computation, including the tuningTableLookup compile path.
• Minor — RockTuningImpl.cpp:1777: no lit coverage for the RockGemmGemmWrapperInterface branch.
• Minor — perfRunner.py:1766: supports_split_k is plumbed two different ways across the five config classes.

Notes

• The extract_tuning_key_metadata call sites look complete: benchmark_conv (:917) and tune_mlir_kernels (:2503) are the two places raw commandline reaches MIOpenDriver, and the hipBLASLt/CK paths go through from_command_line → generate_mlir_driver_commandline, which the new test_split_k_support_metadata_round_trips asserts is clean.
• The comment at RockTuningImpl.cpp:1745-1747 explaining why the func::FuncOp overload is used instead of the ModuleOp one is genuinely useful — it documents a non-obvious hazard.
• Worth double-checking benchmark_fusion_kernels (perfRunner.py:2409-2450): it deliberately looks a fused problem up against an unfused tuning DB and then forces split-K to 1. Making legality part of the key is exactly what makes those two strings diverge.
• The PR description's "rocmlirTriton port of this change:" line has no link.

CI status

C/C++ premerge checks, Python performance script tests, and Python format and lint checks all pass. The only failing entry is this auto-review pipeline's own review check.

[rocmlir-pr-reviewer]

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

---

Scope

Appends -supportsSplitK true|false to the generated tuning key so tuning infrastructure can distinguish otherwise-identical problem keys that map to kernels with different Split-K legality. Touches RockTuningImpl.cpp (key serialization), five lit tests, perfRunner.py (key parse/emit round-trip), and two Python test files.

Findings

The C++ side looks correct: serializeWithSplitKSupport resolves the enclosing func::FuncOp from the op it just serialized rather than the module's first function, which is the right call for multi-function modules, and both the RockGemmWrapperInterface and RockGemmGemmWrapperInterface branches now go through it.

The remaining concern is Python/C++ key parity. PerfConfiguration.supports_split_k defaults to True, but testFusionLegalitySplitK routes through validOutputAtomicAdd (mlir/lib/Dialect/Rock/utility/fusionUtils.cpp:42-52), which fails for any output element type outside f32/f16/bf16 and for archs lacking hasAtomicAdd. So C++ emits false for unfused i8 GEMM and convint8 (i32 output) while canonicalize_config still produces true — an exact-match miss in tuningTableLookup. Details in the thread reply on perfRunner.py.

Test coverage now includes the attention/gemm-gemm path (problem-key.mlir CHECK_1) and three fusability cases covering both polarities, which closes the earlier gaps. Still missing is a check pinning -supportsSplitK false for an unfused problem whose output dtype cannot use atomic add — that is the exact case the Python default gets wrong, so a lit test would make the divergence visible.

Notes

Four of the five threads from the previous review are addressed: the MIGraphX tuning-DB compatibility question, the compile-time cost (measured at ~80 microseconds), the missing gemm-gemm key coverage, and the duplicate supports_split_k plumbing in AttentionConfiguration.

CI status

No failing or cancelled checks. C/C++ premerge checks, Python performance script tests, and Python format and lint checks were still in progress at review time.

[rocmlir-pr-reviewer]

Verdict: COMMENT  ·  New findings: 2 (0 Critical, 1 Major, 1 Minor)

---

Scope

Extends the generated tuning key with -supportsSplitK true|false so tuning infrastructure cannot reuse a cached perf config across two kernels that disagree on split-K legality. On the C++ side, getTuningProblemStr(ModuleOp) now routes both the RockGemmWrapperInterface and RockGemmGemmWrapperInterface branches through a serializeWithSplitKSupport lambda that judges legality on the op's own enclosing func::FuncOp via rock::testFusionLegalitySplitK. On the Python side, perfRunner.py gains extract_tuning_key_metadata, infer_split_k_support, resolve_split_k_support, and lookup_fusion_tuning_config, plus a supports_split_k attribute and tuning_key_metadata() on PerfConfiguration. Lit tests and the Python unit tests are updated accordingly.

Findings

Two non-blocking findings, both in test/infra rather than production logic:

• mlir/utils/performance/tests/mock_hip.py:93 (Major) — the arch mock is arch-insensitive, so the widened default feature set makes every arch look split-K capable and defeats the very sync the new tests are meant to guard.
• mlir/utils/performance/perfRunner.py:2431 (Minor) — original_support is provably always False at that point.

All five issues from the previous review round are addressed: the Python default now derives from the arch/dtype rather than a hardcoded False; the tuning-key break is acceptable given no shipped MIGraphX DBs; the testFusionLegalitySplitK cost was measured at ~80us; the gemm-gemm/attention branch now has --emit-tuning-key coverage; and the AttentionConfiguration constructor now matches its siblings.

Notes

Spot-checked the Python mirror against ground truth: infer_split_k_support's {f32: ATOMIC_ADD, f16: ATOMIC_ADD_F16, bf16: ATOMIC_ADD_BF16} map matches validOutputAtomicAdd in fusionUtils.cpp, and against AmdArchDb.cpp it correctly yields true for f32/f16 on gfx908/gfx90a/gfx942, false for bf16 everywhere except gfx950, and false for i8/i32. The semantics agree; only the test mock diverges.

The MIOpenDriver call sites in benchmark_external and tune_mlir_kernels correctly strip the new metadata before building external argv. Worth a second look at any remaining consumer that splits a to_command_line() string straight into an external tool's argv, since the key now carries a non-rocmlir-gen option.

problem-key.mlir only exercises the new suffix on CHECK_1; extending CHECK_2/CHECK_3 would broaden coverage of the gemm-gemm path.

CI status

No failing or cancelled checks. "Python performance script tests", "C/C++ premerge checks", and "Python format and lint checks" were all still in progress at review time.

[rocmlir-pr-reviewer]

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

---

Scope

Appends -supportsSplitK true|false to generated tuning problem keys so tuning infrastructure never reuses a split-K perf config for a fused kernel that cannot legally use one. C++ side adds a serializeWithSplitKSupport lambda in RockTuningImpl.cpp that calls testFusionLegalitySplitK on the parent func::FuncOp of the op being serialized; Python side (perfRunner.py) parses the new metadata out of argv, infers it from arch/output dtype when absent, and appends it on every to_command_line(). Lit tests and the pytest suite are updated accordingly.

Findings

No blocking issues found. Both threads that were still open on the previous revision are addressed on 7cd05e51:

• mock_hip.py — _mock_lookup_arch_info is now arch-aware, backed by _ARCH_DEFAULT_FEATURES. I cross-checked the table against AmdArchDb.cpp: gfx900 → gcnInfo (no atomics), gfx908/gfx90a/gfx942 → atomic_add | atomic_add_f16, gfx950 → additionally atomic_add_bf16. The affected tests moved off gfx900 onto gfx908, so the Python/C++ key-sync assertions can now actually fail on a mismatch.
• perfRunner.py — lookup_fusion_tuning_config no longer save/restores supports_split_k on the caller's config; it derives the fallback key by swapping the fixed -supportsSplitK false suffix, guarded by an assert.

Notes

Spot-checks that came back clean:

• The base GEMM serializer strips trailing separators before returning, so the leading space in " -supportsSplitK " produces exactly one separator — no double space, and existing keys keep their shape apart from the new suffix.
• raw_svector_ostream is constructed over the existing out buffer, so it appends rather than truncates.
• Using getParentOfType<func::FuncOp>() rather than the ModuleOp overload of testFusionLegalitySplitK is the right call — that overload asserts a single function and inspects the module's first one, which is not necessarily the one holding the serialized op.
• MIOpenDriver argv is sanitized in both benchmarkMIOpen and tune_mlir_kernels, so the non-driver metadata cannot leak into a driver invocation.

Two out-of-scope observations, neither worth blocking on: the mock arch table's NONE default reports RDNA3/RDNA4/gfx1250 as split-K-incapable even though AmdArchDb.cpp gives them atomic_add — harmless today since no test exercises those archs, but worth extending if one is added later. And RockTuningImpl.cpp names func::FuncOp without an explicit FuncOps.h include; it resolves transitively and the C/C++ premerge check is green.

CI status

No genuine CI failures. The only fail entry is this pipeline's own review check, which is expected. Jenkins and its sub-jobs (Build and Test, MIGraphX, Code coverage, parameter sweeps) were still pending at review time; all completed GitHub checks — Python performance script tests, Python format and lint, and C/C++ premerge — passed.

[codecov]

Codecov Report

❌ Patch coverage is 71.42857% with 4 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
mlir/lib/Dialect/Rock/Tuning/RockTuningImpl.cpp	71.43%	2 Missing and 2 partials ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           develop    #2453      +/-   ##
===========================================
+ Coverage    82.57%   83.72%   +1.16%     
===========================================
  Files          120      121       +1     
  Lines        42852    43194     +342     
  Branches      7110     7184      +74     
===========================================
+ Hits         35381    36163     +782     
+ Misses        4815     4493     -322     
+ Partials      2656     2538     -118     

Flag	Coverage Δ
gfx120x	83.63% <71.43%> (+1.10%)	⬆️
gfx950	83.39% <71.43%> (+1.05%)	⬆️
mfma	83.54% <71.43%> (?)

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

Files with missing lines	Coverage Δ
mlir/lib/Dialect/Rock/Tuning/RockTuningImpl.cpp	61.14% <71.43%> (+1.82%)	⬆️

... and 46 files with indirect coverage changes

🚀 New features to boost your workflow:

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