Triton #10646 — Why an Erased Layout Conversion Was Still Reused to Build a Load Address

Layout rematerialization, eager IR deletion, and the side table that continued to remember a value after the operation producing it was gone

The generated IR was valid.

The load address was wrong.

Triton’s RemoveLayoutConversions pass had eliminated one ttg.convert_layout operation and erased it from the IR.

But the pass maintained a separate rematerialization cache across the same walk.

That cache could still lead a later conversion back to the value associated with the erased operation.

The later conversion then reused the stale rematerialization instead of rebuilding the required value in the correct layout.

In the reported PyTorch Inductor failure, the mistake reached pointer arithmetic.

The correct address calculation required a per-row term derived from:

row % 192

The broken rematerialization instead reused a term derived from:

(row / 192) × 128

The resulting program remained structurally legal.

It loaded from the wrong location.

The PR author reports that the miscompiled load propagated to NaNs in the Inductor timm training-accuracy suite. After the repair, seven affected models moved from fail_accuracy to pass on H100.

The complete failure chain was:

ConvertLayoutOp is backward-rematerialized
        ↓
its IR uses are replaced
        ↓
the operation is erased immediately
        ↓
rematerialization state still refers to its result
        ↓
a later conversion requests the same value and encoding
        ↓
the stale entry is accepted as reusable
        ↓
the wrong rematerialized value reaches tt.addptr
        ↓
the kernel loads from the wrong address

PR #10646 did not restore the old delayed-cleanup system.

It made the eager deletion safe.

Before erasing the conversion, the patch routes its result through the same replacement bookkeeping used for other deleted values.

That lets Triton repair the side tables first.

Only then is the operation removed.

Deleting an IR operation is safe only after every analysis structure has stopped treating its result as reusable state.

Why Triton tries to remove layout conversions

A Triton tensor layout describes how logical tensor elements are distributed across GPU execution resources.

Depending on the layout, elements may be assigned differently across:

threads,
warps,
register lanes,
or CTA-level structures.

A ttg.convert_layout changes that physical distribution.

Conceptually:

logical tensor values

layout A
        ↓
ttg.convert_layout
        ↓
same logical values

layout B

Some conversions require real data movement between threads.

That can be expensive.

But a conversion is not always necessary.

Suppose a value was produced by a cheap operation:

broadcast

arithmetic

comparison

index calculation

Instead of:

produce in layout A

then

convert to layout B

the compiler may be able to recompute the producer directly in layout B.

rematerialize producer in layout B

→ no separate conversion

This is backward rematerialization.

The pass starts from a conversion and walks backward through operations that can be reproduced safely.

It then clones that slice using the target encoding.

Old path:

producer in A
→ convert A to B
→ consumer

Rematerialized path:

producer rebuilt in B
→ consumer

The optimization can remove data movement while preserving the same logical computation.

The bug was not in the idea of rematerialization.

It was in the lifetime of the state used to reuse previous rematerializations.

Why the pass keeps a rematerialization cache

One function may contain several conversions requesting the same value in the same target layout.

Rebuilding the same backward slice every time would duplicate work.

At the affected revision, LayoutRematerialization therefore kept state equivalent to:

rematMapping:

(original value, requested encoding)
→ rematerialized value

It also tracked which original values had already been rematerialized.

Conceptually:

Original value V
+
target layout B
        ↓
rematerialized value V_B

Later in the same pass:

another conversion asks for:

V in layout B
        ↓
cache lookup finds V_B
        ↓
reuse V_B

That reuse is correct only while V_B remains a live and semantically valid value.

The map is not merely a performance cache.

It participates in compiler correctness.

Every stored value must continue to satisfy three conditions:

1. It still exists in the IR.

2. It still represents the original value.

3. It still has the requested encoding.

An erased operation can violate the first condition.

An incorrectly redirected mapping can violate the second or third.

The earlier cleanup change altered the lifetime boundary

Triton PR #9511 changed when dead rematerialization operations were erased.

Before that cleanup, dead operations were collected in a set and removed after the pass had finished visiting all conversions.

During walk:

rewrite conversions

retain dead operations temporarily

keep remat state valid

After walk:

erase collected dead operations

PR #9511 moved to eager cleanup.

The rewritten operations were erased as soon as their replacement work completed.

rewrite one conversion

update replacement state

erase dead operations immediately

continue walk

The motivation was reasonable.

Dead rematerialized operations could be temporarily malformed—for example, a structured operation whose region had already been replaced.

Keeping such operations in the IR until the end made later manipulation more fragile.

The eager-cleanup change assumed that updateRematMapping() removed or redirected every side-table entry associated with the operations being deleted.

Most deleted values followed that path.

The root ConvertLayoutOp did not.

The root conversion bypassed the replacement list

Inside rewriteSlice(), Triton already accumulated replacements for values produced by operations that were being rewritten or deleted.

Conceptually:

replacements:

old value A → new value A'

old value B → new value B'

old value C → new value C'

The pass then called:

updateRematMapping(replacements);

That function migrated or removed entries associated with the old values.

After the side tables were updated, the pass replaced the IR uses and erased the dead operations.

The rewritten root conversion was handled differently.

The affected code did this:

convertOp.replaceAllUsesWith(
    mapping.lookup(convertOp.getSrc())
);

updateRematMapping(replacements);

convertOp->erase();

The conversion’s uses changed.

But its result was never added to:

replacements

Therefore:

IR uses
→ updated

while:

rematMapping / mappedValues
→ not told that the convert result was being deleted

The operation was then erased.

The side table survived the object whose lifetime it was supposed to track.

The minimum two-conversion scenario

The failure required more than one conversion in the same rematerialization walk.

A simplified shape is:

source value
        ↓
broadcast A
        ↓
convert A
        ↓
consumer A

and later:

same source value
        ↓
broadcast B
        ↓
convert B
        ↓
pointer calculation

The first conversion is eliminated through backward rematerialization.

broadcast A is rebuilt in the target layout

convert A becomes unnecessary

Its uses are replaced.

The conversion is erased.

But the cache still contains state associated with the eliminated conversion or the rematerialized path that led through it.

The second conversion then asks:

Do we already have this value
in the target encoding?

The old state answers:

yes

That answer should have been invalidated or redirected when the first conversion disappeared.

Instead, the pass reuses it.

First conversion:

creates reusable remat state
        ↓
conversion erased
        ↓
state not cleaned

Second conversion:

queries same value/layout pair
        ↓
receives stale remat result

A function containing only one relevant conversion would not expose this.

The pass had to eliminate one conversion and then continue far enough for another conversion to consult the retained state.

The real address expressions were not interchangeable

The reported Inductor case involved two row-derived values.

The address needed a term based on:

row % 192

The stale rematerialization substituted a term derived from:

(row / 192) × 128

These expressions differ immediately after the first 192-row block.

Take:

row = 193

The intended term is:

The reused division-based term is:

(193 / 192) × 128

=

1 × 128

=

128

The difference is not a minor rounding error.

correct offset component:
1

stale component:
128

Once that value reaches:

tt.addptr

the pointer remains a well-typed pointer.

It simply refers to another location.

The compiler can therefore produce valid IR and valid machine code for the wrong address calculation.

Why the verifier did not reject the result

Compiler verification can prove many structural properties.

For example:

Does every operand dominate its use?

Do tensor ranks and element types match?

Does the pointer arithmetic consume an integer offset?

Does each operation satisfy its declared type contract?

The wrong rematerialized value could satisfy all of them.

Both candidate expressions produced compatible integer tensors.

row % 192
→ integer tensor

(row / 192) × 128
→ integer tensor

Both could be broadcast to the required shape.

Both could be added to a pointer.

The missing property was semantic provenance:

Is this the specific row term
that the original computation required?

Ordinary IR verification does not reconstruct that mathematical intent.

The final program therefore did not need to crash during compilation.

It could pass verification, lower successfully, and run with a silently altered address.

The fix made the root conversion an ordinary replacement

PR #10646 changed two effective lines of production logic.

Instead of replacing the conversion’s uses through a special path:

convertOp.replaceAllUsesWith(
    mapping.lookup(convertOp.getSrc())
);

the patch records:

replacements.emplace_back(
    convertOp.getResult(),
    mapping.lookup(convertOp.getSrc())
);

The conversion now enters the same replacement sequence as the other values removed by rewriteSlice().

convert result
        ↓
added to replacements
        ↓
updateRematMapping()
        ↓
side-table entries removed or redirected
        ↓
IR uses replaced
        ↓
conversion erased

This ordering matters.

The patch does not erase the operation and then attempt to repair dangling analysis state.

It identifies the replacement while both values are still available.

The mapping update runs before deletion.

Current Triton main retains this sequence: the root conversion is added to replacements, updateRematMapping() runs, all uses are replaced through the shared loop, and only then is convertOp erased.

Why the patch did not return to delayed deletion

The regression was introduced after eager deletion changed the relationship between IR lifetime and rematerialization-state lifetime.

One possible repair would have been:

restore the old end-of-pass deletion queue

PR #10646 chose a narrower correction.

keep eager deletion

but

make the bookkeeping complete

That preserves the original cleanup improvement from #9511:

dead operations do not remain in the IR for the rest of the walk;
temporarily malformed rewritten operations are removed promptly;
later transformations do not encounter unnecessary dead nodes.

The missing rule is added exactly where the root conversion crosses the deletion boundary.

Every deleted value affecting remat state
must enter updateRematMapping()

The eager-cleanup design remains.

Its side-table contract becomes explicit.

The final regression was intentionally minimized

The original failure came from an Inductor Welford-style kernel and a detailed load-address chain.

During review, maintainers asked for the MLIR test to be reduced.

The test was repeatedly trimmed until the checked-in form contained the minimum repeated-rematerialization structure rather than the entire timm address expression.

The final test contains:

one source tensor

two broadcasts of that source

two conversions into the same layout

one converted value used by a comparison

one converted value used by tt.addptr

a load and a loop-carried computation

After running:

-tritongpu-remove-layout-conversions

FileCheck requires:

two tt.broadcast operations

exactly one remaining ttg.convert_layout

The regression is no longer a literal copy of:

row % 192

or:

(row / 192) × 128

Its job is narrower.

It preserves the state-lifetime condition that made the real address substitution possible:

one conversion is eliminated

the walk continues

a later equivalent rematerialization request occurs

both required producer paths must remain distinct and valid

The minimized test and the production guard both remain present in current main.

Runtime validation reached real training models

The checked-in regression is a compiler-level MLIR test.

The PR also reports end-to-end validation on:

NVIDIA H100

PyTorch Inductor

inductor_timm --accuracy --training --amp

The affected models were:

inception_v3

adv_inception_v3

ghostnet_100

repvgg_a2

tf_efficientnet_b0

mobilenetv2_100

mobilenetv3_large_100

Before the corrected load-address rematerialization, those models were reported as:

fail_accuracy

After the fix:

pass

This is stronger than a theoretical stale-map concern.

The author connected the pass defect to a user-visible training-accuracy failure and validated the corrected pipeline on hardware.

It should still be described precisely.

Checked-in evidence:
MLIR transformation regression

Runtime evidence:
author-reported H100 Inductor validation

The hardware run was not independently reproduced for this article.

Why the failure could look unrelated to layout conversion

The final symptom was:

NaN during model training

That does not immediately suggest:

an erased ConvertLayoutOp remained in a rematerialization side table

The path crossed several abstraction layers.

layout-conversion cleanup
        ↓
stale compiler analysis state
        ↓
wrong value selected during rematerialization
        ↓
wrong pointer offset
        ↓
wrong tensor element loaded
        ↓
invalid numerical state
        ↓
accuracy failure

The conversion itself did not necessarily carry model data into the wrong layout at runtime.

It influenced which address-producing computation the compiler rebuilt.

This is why the most useful regression sits near the compiler pass rather than only at the model level.

The model test detects the outcome.

The MLIR test protects the cause.

What the merged evidence confirms

The reviewed source and PR record confirm that:

LayoutRematerialization reused previously generated values across later conversions;
PR #9511 changed dead-operation cleanup from deferred to eager;
the root rewritten ConvertLayoutOp was erased without being added to the replacement list;
updateRematMapping() therefore did not process that conversion result;
a later conversion could reuse stale rematerialization state;
the reported Inductor case miscompiled a load-address expression;
the PR author observed NaNs and timm training-accuracy failures;
the fix adds the conversion result to replacements before updating the maps and erasing the operation;
a minimized MLIR regression protects repeated rematerialization during one walk;
and both the fix and regression remain in current Triton main.

What the evidence does not establish

The public record does not establish that:

every removed layout conversion was affected;
every stale entry produced a wrong load address;
the bad pointer necessarily accessed memory outside its allocation;
a memory-safety exploit was demonstrated;
all GPUs or Triton backends reproduced the same model failure;
the fix changes kernel performance;
or every possible side table in the pass is now formally proven coherent.

The supported conclusion is narrower:

Eager erasure of a backward-rematerialized layout conversion could leave its result represented in Triton’s rematerialization state. A later conversion could reuse that stale state and build a valid but semantically incorrect computation, including a wrong load address.

The contract that remains

Compiler state does not live only in the IR.

A transformation may also keep:

value maps,
alias information,
dominance summaries,
rematerialization caches,
replacement lists,
and other side tables.

Erasing one operation from the visible graph does not erase it from those structures automatically.

The safe order is:

identify replacement

        ↓

update every side table

        ↓

replace IR uses

        ↓

erase old operation

Not:

erase operation

        ↓

assume every analysis forgot it

An IR value is not dead while any compiler decision can still retrieve it as live state.

That is the lifetime contract repaired by Triton #10646.

Related material

Patch status: Merged into Triton main on June 23, 2026, as 3bb23dde; the bookkeeping fix and minimized regression remain in current main.

Search This Blog

Resonetic Lab | AI Systems Code Review