[AIROCMLIR-1246] Ignore negligible anchors when resolving layout conflicts

[dhernandez0]

Motivation

rock-narrow-redundant-loads rewrites a broadcast-invariant load (e.g. a 128x64 bias tile) into a 128x1 load plus a broadcast. That narrow load still anchors a layout in remove-layout-conversions, and since no candidate is an mma layout yet before accelerate-matmul, iteration order handed the dot's whole epilogue the load's sizePerThread = [1, 1] layout plus a full-tile shared-memory round trip: up to 1.6x slower on gfx1100 convolution+reduce fusions.

Technical Details

Adds the tie-break LayoutPropagation::resolveConflicts asked for in its "hacky resolve" comment. Each anchor carries the traffic its layout generates, the figure propagates with the layout, and a candidate whose anchor generates no more than one dword per thread loses to one that governs more. The existing type-based preference still decides first, so this only settles what it leaves tied.

Perf numbers

Results of develop vs this branch:

Brach	best config time
develop	5.056 ms
develop, narrowing pass disabled	4.793 ms
this branch	4.703 ms

module {
  func.func @mlir_reshape_convolution_reshape_broadcast_add_convert_mul_reshape_reduce_sum_reshape_mul_mul_reshape_reduce_sum_reshape(%arg0: !migraphx.shaped<1x32x4x1024x1024xf16, 134217728x4194304x1048576x1024x1>, %arg1: !migraphx.shaped<128x128x3x3xf16, 1152x9x3x1>, %arg2: !migraphx.shaped<32x4xf16, 4x1>) -> (!migraphx.shaped<1x32x1x1x1xf32, 32x1x1x1x1>, !migraphx.shaped<1x32x1x1x1xf32, 32x1x1x1x1>, !migraphx.shaped<1x32x4x1024x1024xf16, 134217728x4194304x1048576x1024x1>) attributes {rock.arch = "gfx1100", rock.enable_splitk_for_tuning, rock.kernel} {
    %0 = migraphx.literal(dense<2.38418579E-7> : tensor<1xf32>) : <1xf32, 0>
    %1 = migraphx.literal(dense<2.38418579E-7> : tensor<1xf32>) : <1xf32, 0>
    %2 = migraphx.reshape %arg0 {dims = [1, 128, 1024, 1024]} : <1x32x4x1024x1024xf16, 134217728x4194304x1048576x1024x1> -> <1x128x1024x1024xf16, 134217728x1048576x1024x1>
    %3 = migraphx.convolution %2, %arg1 {dilation = [1, 1], group = 1 : i64, padding = [1, 1, 1, 1], padding_mode = 0 : i64, stride = [1, 1]} : <1x128x1024x1024xf16, 134217728x1048576x1024x1>, <128x128x3x3xf16, 1152x9x3x1> -> <1x128x1024x1024xf16, 134217728x1048576x1024x1>
    %4 = migraphx.reshape %3 {dims = [1, 32, 4, 1024, 1024]} : <1x128x1024x1024xf16, 134217728x1048576x1024x1> -> <1x32x4x1024x1024xf16, 134217728x4194304x1048576x1024x1>
    %5 = migraphx.broadcast %arg2 {axis = 1 : i64, out_dyn_dims = [], out_lens = [1, 32, 4, 1024, 1024]} : <32x4xf16, 4x1> -> <1x32x4x1024x1024xf16, 0x4x1x0x0>
    %6 = migraphx.add %4, %5 : <1x32x4x1024x1024xf16, 134217728x4194304x1048576x1024x1>, <1x32x4x1024x1024xf16, 0x4x1x0x0> -> <1x32x4x1024x1024xf16, 134217728x4194304x1048576x1024x1>
    %7 = migraphx.convert %6 {target_type = 2 : i64} : <1x32x4x1024x1024xf16, 134217728x4194304x1048576x1024x1> to <1x32x4x1024x1024xf32, 134217728x4194304x1048576x1024x1>
    %8 = migraphx.multibroadcast %1 {out_dyn_dims = [], out_lens = [1, 32, 4, 1024, 1024]} : <1xf32, 0> -> <1x32x4x1024x1024xf32, 0x0x0x0x0>
    %9 = migraphx.mul %7, %8 : <1x32x4x1024x1024xf32, 134217728x4194304x1048576x1024x1>, <1x32x4x1024x1024xf32, 0x0x0x0x0> -> <1x32x4x1024x1024xf32, 134217728x4194304x1048576x1024x1>
    %10 = migraphx.reshape %9 {dims = [1, 32, 4194304]} : <1x32x4x1024x1024xf32, 134217728x4194304x1048576x1024x1> -> <1x32x4194304xf32, 134217728x4194304x1>
    %11 = migraphx.reduce_sum %10 {axes = [2]} : <1x32x4194304xf32, 134217728x4194304x1> -> <1x32x1xf32, 32x1x1>
    %12 = migraphx.reshape %11 {dims = [1, 32, 1, 1, 1]} : <1x32x1xf32, 32x1x1> -> <1x32x1x1x1xf32, 32x1x1x1x1>
    %13 = migraphx.multibroadcast %0 {out_dyn_dims = [], out_lens = [1, 32, 4, 1024, 1024]} : <1xf32, 0> -> <1x32x4x1024x1024xf32, 0x0x0x0x0>
    %14 = migraphx.mul %7, %7 : <1x32x4x1024x1024xf32, 134217728x4194304x1048576x1024x1>, <1x32x4x1024x1024xf32, 134217728x4194304x1048576x1024x1> -> <1x32x4x1024x1024xf32, 134217728x4194304x1048576x1024x1>
    %15 = migraphx.mul %14, %13 : <1x32x4x1024x1024xf32, 134217728x4194304x1048576x1024x1>, <1x32x4x1024x1024xf32, 0x0x0x0x0> -> <1x32x4x1024x1024xf32, 134217728x4194304x1048576x1024x1>
    %16 = migraphx.reshape %15 {dims = [1, 32, 4194304]} : <1x32x4x1024x1024xf32, 134217728x4194304x1048576x1024x1> -> <1x32x4194304xf32, 134217728x4194304x1>
    %17 = migraphx.reduce_sum %16 {axes = [2]} : <1x32x4194304xf32, 134217728x4194304x1> -> <1x32x1xf32, 32x1x1>
    %18 = migraphx.reshape %17 {dims = [1, 32, 1, 1, 1]} : <1x32x1xf32, 32x1x1> -> <1x32x1x1x1xf32, 32x1x1x1x1>
    return %12, %18, %6 : !migraphx.shaped<1x32x1x1x1xf32, 32x1x1x1x1>, !migraphx.shaped<1x32x1x1x1xf32, 32x1x1x1x1>, !migraphx.shaped<1x32x4x1024x1024xf16, 134217728x4194304x1048576x1024x1>
  }
}

Test Plan

PR CI.

Test Result

All tests pass.

Submission Checklist

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

[Copilot]

🟡 Changes recommended

Minor but actionable correctness/maintainability issues were found in the updated Triton transform implementation (see PR comments).

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

This PR improves Triton’s tritongpu-remove-layout-conversions conflict resolution by adding a traffic-based tie-breaker so “negligible” layout anchors (e.g., narrow 128x1 loads that don’t meaningfully describe downstream tile layouts) don’t inadvertently force expensive shared-memory round trips in dot epilogues. It also adds targeted regression tests to cover both the Triton-side heuristic and the rocMLIR↔Triton interaction with Rock’s load-narrowing pass.

Changes:

• Track per-encoding “anchor traffic” during layout propagation and use it as a tie-break in LayoutPropagation::resolveConflicts.
• Add a Triton GPU dialect test validating that narrow-load anchors don’t win layout conflicts when their traffic is negligible.
• Add a rocMLIR Rock integration test ensuring narrowed loads remain narrow through Triton’s layout passes without contaminating epilogue layouts.

File summaries

File	Description
triton-patches/triton-patch-content.txt	Registers the new downstream Triton patch and documents its motivation/tests.
triton-patches/patch-negligible-anchor-layout-conflict.patch	Downstream patch recording the Triton transform + test additions.
external/triton/lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp	Implements traffic propagation and the negligible-anchor tie-break in conflict resolution.
external/triton/test/TritonGPU/amd/remove-layout-conversions-negligible-anchor.mlir	New regression test for the negligible-anchor conflict behavior.
mlir/test/Dialect/Rock/narrow-redundant-loads-through-triton.mlir	New end-to-end test covering Rock narrowing interacting with Triton layout passes.

Review details

• Files reviewed: 5/5 changed files
• Comments generated: 2
• Review effort level: Lite

---

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

[dhernandez0]

fixed

[dhernandez0]

fixed

[rocmlir-pr-reviewer]

This new lit test does not appear to be run by any rocmlirTriton CI target. external/triton/CMakeLists.txt:543 guards add_subdirectory(test) behind TRITON_BUILD_UT, and cmake/triton.cmake:154 sets TRITON_BUILD_UT OFF; cmake/triton.cmake:249-266 then adds only ${TRITON_PROJECT_DIR}/test/lib, not test/ itself. So the external/triton/test/ lit suite is never configured here and this file is dead coverage in this repo, even though the behavior it pins is exactly what the PR changes. Please either confirm how it gets executed in our CI, or move/duplicate the two cases under mlir/test/ where check-rocmlir picks them up. This maps to the checklist's Major bullet on a pass change without both positive and negative Lit coverage that actually runs.

[dhernandez0]

we don't run triton CI in our CI. That's is a known issue that needs to be fixed in another PR. I've verified the new triton tests pass manually.

[rocmlir-pr-reviewer]

getAnchorTraffic relies on getByteCount, which leaves dtypeBitWidth at 0 unless the element type satisfies isIntOrFloat(). A tensor<...x!tt.ptr<T>> fails that check, so every pointer-tensor anchor scores 0 bytes and is therefore unconditionally classified negligible, regardless of its real size. That matters here because initAnchorLayout explicitly anchors function arguments (lines 288-293), and a pointer-tensor argument such as tensor<128x64x!tt.ptr<f32>> will now always lose the tie-break to any other candidate — silently inverting the intended preference on pointer-heavy kernels. Suggest special-casing triton::PointerType in getAnchorTraffic (e.g. use the pointee's bit width, or element count times the access width) so pointer tensors are scored by the data they actually move. This is the checklist's "silently falls through for unhandled types" Major case.

[dhernandez0]

fixed

[rocmlir-pr-reviewer]

The --implicit-check-not='convert_layout {{.*}}tensor<128x64xf32' guard only covers f32 tiles, but the epilogue this test is protecting is f16 — the CHECK lines at 32-34 match tensor<128x64xf16 operations. A regression that reintroduces a full-tile convert_layout on the f16 add/broadcast would pass this test unnoticed. Consider broadening the guard to convert_layout {{.*}}tensor<128x64x so it covers both element types, after confirming no legitimate 128x64 conversion survives in the printed IR for this config.

[rocmlir-pr-reviewer]

Verdict: COMMENT  ·  Findings: 4 (0 Critical, 2 Major, 2 Minor)

---

Scope

Adds a traffic-based tie-break to Triton's LayoutPropagation::resolveConflicts so that an anchor whose layout governs no more than one dword per thread loses to one governing more. Motivated by rock-narrow-redundant-loads rewriting a broadcast-invariant bias tile into a 128x1 load, which then anchored the dot epilogue's layout and cost up to 1.6x on gfx1100. Touches the vendored Triton pass ([EXTERNAL] commit), its downstream patch record, and two new lit tests.

Findings

• external/triton/test/TritonGPU/amd/remove-layout-conversions-negligible-anchor.mlir:1 (Major) — the new Triton-level lit test does not appear to be executed by any rocmlirTriton CI target.
• external/triton/lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp:64 (Major) — getByteCount returns 0 for pointer-element tensors, so any pointer-tensor anchor is unconditionally scored negligible.
• mlir/test/Dialect/Rock/narrow-redundant-loads-through-triton.mlir:25 (Minor) — the --implicit-check-not guard only covers f32 tiles while the epilogue it protects is f16.
• triton-patches/triton-patch-content.txt:356 (Minor) — the new patch record omits the "re-evaluate / drop once upstreamed" exit criterion its neighbours carry.

Notes

Verified and clean: the vendored-tree edits are isolated in a correctly [EXTERNAL]-prefixed commit; triton-patches/patch-negligible-anchor-layout-conflict.patch matches that commit's diff exactly and is indexed in triton-patch-content.txt; TritonGPUDialect::getThreadsPerWarp(ModuleOp) is declared in the vendored TritonGPUDialect.td, so the new call compiles; the rewritten selection loop preserves the previous "first encoding wins" fallback, and the propagation fixpoint still terminates since traffic is monotone and bounded.

The rocMLIR back-port check does not fire — the touched paths are confined to external/triton/, triton-patches/, and mlir/test/, none of which are in the shared path list.

Neither new .mlir file carries a license header, but sibling first-party tests (mlir/test/Dialect/Rock/narrow-redundant-loads.mlir, triton-to-hsaco-denormal-mode.mlir) do not either, so this reads as established repo practice rather than a defect this PR introduces; worth settling repo-wide instead of here.

Perf-wise, the extra SmallDenseMap<Attribute, int64_t, 8> per LayoutInfo plus traffic changes feeding hasChanged can add propagation iterations and memory on large kernels. A compile-time spot-check on a big fusion would be reassuring.

CI status

No checks are in a fail or cancel state. Jenkins, "Build and Test", MIGraphX, and ml-ci-internal.amd.com are still pending, so the change has not yet been validated by the full suite.

[pabloantoniom]

Could you double-check that this PR does not lose the performance gains achieved by #445 and/or the original NarrowLoad PR? I just want to make sure that we are not losing performance by changing the layout in unexpected ways.

Also, would be good to check on gfx942/gfx950 as suggested in my comment

Lastly, I don't know the code well but I assume you will upstream this to Triton so they should provide proper feedback on that. The fact that we are replacing Hacky resolve and TODO: add a proper heuristic with a proper heuristic makes me think the should be interested in something like this (hopefully)

[pabloantoniom]

Can you explain the reason for this heuristic? Assume:

lookupNumWarps = 8
getThreadsPerWarp = 64 (gfx942)

Then negligibleTraffic = 2048 bytes. Meaning that tensor<32x32xf16> would count as negligible. Not sure if that would work well?

[dhernandez0]

sure, the idea is that a tensor is negligible if it's a single register per thread. In this case, there are 8*64=512 threads, so a 32x32 f16 tensor would fit into 1 register per thread.

The existing isExpensiveLoadOrStore has a similar heuristic. Just without taking into account datatypes and register size.

[pabloantoniom]

This may work well for gfx1101 but did you test on gfx942/gfx950?

[dhernandez0]

I'll check this, but the logic is the same IMO.

[pabloantoniom]

The threshold is different depending on threadsPerWarp so I'm not sure if it needs some adjustment to take that into account

[justinrosner]

Just a few small questions. Also some of the AI review comments look valid

[dhernandez0]

Could you double-check that this PR does not lose the performance gains achieved by #445 and/or the original NarrowLoad PR? I just want to make sure that we are not losing performance by changing the layout in unexpected ways.

Checked, perf stays flat for the mlir that made us open that ticket.

Lastly, I don't know the code well but I assume you will upstream this to Triton so they should provide proper feedback on that. The fact that we are replacing Hacky resolve and TODO: add a proper heuristic with a proper heuristic makes me think the should be interested in something like this (hopefully)

Yes, I'll open an upstream PR.

[justinrosner]

Changes look mostly good to me now. Just the one comment. Also make sure that you update the commit structure when it's time to merge in (i.e., squashing to two commits)

[justinrosner]

!encoding now has double meaning as both "no candidate chosen yet" and "the chosen candidate is null", and since rank() now runs on every element instead of breaking at the first preferred-kind match, a null Attribute (which I believe initAnchorLayout can insert) would assert inside isa<BlockedEncodingAttr> and rank itself above real candidates?

[dhernandez0]

fixed

[justinrosner]

Changes look good to me now. It would be good to also run the Triton LIT test suite to make sure that no other tests need to be updated now that we will soon be enabling those in the Nightly CI (see here: #494)