[Nvidia][TritonGPU] Add phases to RemoveLayoutConversions

[zeroherolin]

Background

A ttg.convert_layout (cvt) usually lowers to a cross-thread data reshuffle that round-trips through shared memory, making it one of the more expensive ops in TritonGPU. Yet tensor layouts are chosen independently by several upstream passes: Coalesce picks coalescing-friendly blocked layouts for memory ops, AccelerateMatmul picks MMA layouts for dots, and so on. These decisions are uncoordinated, so cvts are inserted at the layout boundaries to bridge them.

The tritongpu-remove-layout-conversions (RLC) pass, which is responsible for cleaning up redundant cvts, natively only performs forward propagation from anchors and resolves conflicts with a fixed heuristic ("prefer blocked for memory ops, mma otherwise"). This has two structural limitations:

• No backward path: the layout a downstream store/atomic prefers cannot flow back to upstream producers, so residual cvts remain on write-back chains (RNG, mask, reduction results, etc.).
• No global cost view: the fixed rule only looks at a single op's type and never weighs a layout's actual data-movement cost along the whole chain, so small components that could be eliminated as a whole get split apart by the hard rule. The cvt on an MMA-result write-back is especially stuck, because its source (the dot accumulator) is an anchor that neither forward nor backward remat can touch.

The result is that many kernels retain redundant cvts on the hot path — each a shared-memory round-trip — hurting bandwidth on memory-bound kernels and the epilogue on compute-bound ones.

Motivation: without changing RLC's existing semantics and with zero impact on other backends, add four capabilities — backward preference propagation, global cost-based selection, whole small-component solving, and MMA write-back rematerialization — to systematically remove these residual cvts.

Design

This PR non-intrusively enhances the tritongpu-remove-layout-conversions (RLC) pass with four new optimization phases, eliminating redundant ttg.convert_layout (cvt) ops produced by the independent layout decisions of upstream passes (Coalesce, AccelerateMatmul, etc.). The enhancement is enabled by default only on the Nvidia backend (transparent to users, no configuration required) and leaves other backends untouched. To turn it off at runtime, set the environment variable FLAGTREE_RLC_ENHANCE=0, which reverts to the original pass. In addition, all enhancement code is wrapped in the compile-time macro __FLAGTREE_RLC_ENHANCE__ (enabled by default for Nvidia builds); undefining it in CMake removes the enhancement and reverts the pass to behavior identical to the upstream implementation, serving as an absolute fallback if issues arise. All four phases are TLE-aware: they skip TLE cluster remote-address (tle.remote_pointers) chains, leaving their layout intact.

• Phase 1a — Cost-based conflict resolution: When a value has multiple candidate encodings, estimate the data-movement cost of each candidate across all uses as 32 × byteCount and pick the minimum, replacing the hard rule of "prefer blocked for memory ops, mma otherwise"; on ties, fall back to the legacy rule to guarantee zero regression.
• Phase 1b — Local backward propagation: Native RLC only propagates layouts forward, so the layout preferences of downstream stores/atomics never reach upstream producers. For residual write-back cvts that pass a whitelist (blocked→blocked, single-use on both sides, tail chain reaching a store/atomic, producer chain free of reductions), inject the result-side encoding backward as a candidate, then let forward propagation spread it and Phase 1a resolve uniformly. This covers RNG/sampling write-back chains, mask computation chains, etc.
• Phase 2 — Local small-component solving: For closed small components such as reduction write-back chains, loop boundaries, and producer-side chains, simulate the rewrite via a proposal mechanism and count the cvt delta. The proposal is committed as a whole only if all conditions hold — no cvt increase, no memory-access degradation (vetoes for scatter amplification, de-coalescing, reduction/scan crossing, etc.), and a strict loop-depth-weighted cost benefit. The proposal is collected to stay type-consistent (loop-carried values and both scf.if branches are retagged together); otherwise the original IR is kept intact.
• Phase 3 — Store layout rematerialization: For global write-backs (tt.store / result-unused tt.atomic_rmw) whose stored value comes from a single-use cvt with an MMA-layout source, recompute the store's ptr/mask pure-index chains in the value's layout — without degrading coalescing — thereby removing the cvt from the hot data path. Includes a ttg.local_store counterpart and access-preserving duplication of small-vector loads.

Performance

Performance gains are observed for 100+ Ops. Representative gains (median):

Op	dtype	Representative shapes	Δ% (over profitable shapes)	cvt
replication_pad3d	bf16/fp16	1x1x64x128x128 ~ 2x32x16x128x128 (5/5 shapes)	+26.2% ~ +71.0%	5→0
mm / mm_out	fp16	large GEMMs, M/N 4k~152k (49/121 shapes)	+1.2% ~ +53.3% (single point 229.7ms→143.0ms)	13→11
col2im	bf16/fp16	2x64 / 4x32 large-window combos	+36.7% ~ +48.0%	27→0
batch_norm	bf16/fp16	hotspot 20x6x65536	+2.8% ~ +41.7%	18→0
max_pool2d_backward	bf16/fp16	4x3x224x224 etc. (bf16 5/5)	+0.8% ~ +37.9%	90→0
fft	complex64	128x512 ~ 4096x1024	+18.1% ~ +36.4%	80→4
masked_select	bf16/fp16	4096x4096 ~ 10000x65536	+7.9% ~ +30.4%	15→0
router_gemm	bf16	(16~1024)x8x4096 (decode)	+14.6% ~ +26.1%	7→0
topk	bf16/fp16	up to 128x32768x256	+1.7% ~ +21.2%	30→4
RNG family (rand/randn/uniform_/bernoulli_/dropout)	bf16/fp16	4096x4096 ~ 1G elements (all shapes)	+1% ~ +16%	8→0 / 32→0 etc.
Reduction family (amax/max/min/all/any/std/prod etc.)	bf16/fp16	hotspot 1024x1024x1024	+1% ~ +13%	9→5 / 17→3 etc.
fp8_einsum	fp8_e4m3	small M: (1~32)x16x7168x1024	+0.8% ~ +9.0%	0→0 (IR structural optimization)
sort / argsort	bf16/fp16	1024x1024 ~ 4096x4096	+1.6% ~ +8.5%	41→2

Net cvt elimination is ~68% ~ 79% overall. A very small number of ops show minor regressions on individual shapes; all of them stem from the same source as larger gains within the same kernel (inseparable at compile time), have been individually attributed and verified, and are an acceptable trade-off. The overall net result is strongly positive.

[CLAassistant]

All committers have signed the CLA.

[zeroherolin]

TLE performance is unaffected. A/B on Hopper (master switch ON vs OFF) shows the compute-bound TLE kernels (sparse-MLA TLE / FlashMLA-Prefill / Pipelined across all SKV sizes) within <0.5% (measurement noise), and the distributed cluster-GEMM within ~1% noise. This is expected: the phases either don't touch TLE anchors (dot/wgmma/TMA/pipeline) or explicitly skip TLE cluster remote-address chains, so TLE codegen is essentially unchanged.

[zhzhcookie]

LGTM