Intel Triton XPU #7791 — Why a Boundary-Shaped Mask Was Not Enough to Prove It Removable

Intel Triton XPU RemoveMasks, Part 2 of 3 — Canonical loop bounds, mismatched N and END, dynamic dimensions, and equality predicates that did not imply an all-lanes condition

Part 1 examined a crash in Intel Triton’s TritonIntelRemoveMasks pass.

The pass proved that a masked load would never access memory.

But the load had no explicit other operand.

The old rewrite attempted to replace the load result with that absent operand:

Always-false mask
        ↓
`other` does not exist
        ↓
replace uses with a null Value
        ↓
compiler crash

PR #7791 repaired that path by materializing a typed zero.

The same patch also fixed a different and more dangerous class of error.

The compiler could remove a mask even though the condition guarding the new fast path did not actually prove that the mask was true.

The problem appeared in three forms:

A loop upper bound had the expected arithmetic shape,
but it used a different `N` or tile width from the mask.
The loop upper bound was constant,
but the mask’s `N` was still dynamic.
A comparison resembled a boundary check,
but its predicate was `eq` or `ne`
rather than an ordered range predicate.

All three failures came from the same mistaken substitution:

The expression looks like a known mask pattern.

was treated as though it meant:

The generated scalar condition logically implies
that every lane of this mask is true.

Those statements are not equivalent.

A removable mask needs a proof of implication, not merely a familiar expression shape.

PR #7791 was merged with targeted MLIR regressions for all six reported defects, including the three proof-boundary corrections discussed here.


What loop versioning is trying to do

Consider a tiled loop processing N elements in groups of END.

tile width:
END
number of loop iterations:
ceil(N / END)

A simplified Triton-style loop is:

for i in 0 .. ceil(N / END):

    lane =
        range(0, END)

    offset =
        i * END + lane

    mask =
        offset < N

    value =
        load(pointer + offset, mask)

The mask is needed because the last tile may be partial.

For example:

N = 70
END = 32

The iterations cover:

iteration 0:
offsets 0–31
→ all valid

iteration 1:
offsets 32–63
→ all valid

iteration 2:
offsets 64–95
→ only 64–69 valid

The final tile needs a mask.

But when:

N = 64
END = 32

both tiles are complete.

iteration 0:
0–31
→ all valid

iteration 1:
32–63
→ all valid

A compiler can create two versions of the loop.

if the full-tile condition is true:

    run a fast loop
    with unmasked loads

else:

    run the original loop
    with masks retained

Conceptually:

if N is divisible by END:

    unmasked loop

else:

    masked loop

The actual pass also imposes an additional size condition, but the central idea is unchanged.

Loop versioning moves the expensive mask out of the repeated operation and into one scalar branch.


The fast-path condition must imply the vector mask

The optimization is correct only if the following relationship holds:

Versioning condition
        ⇒
every lane of every removed mask is true

More formally, for each operation whose mask is removed:

V(N, END)
        ⇒
M(i, lane, N, END)

for:

every loop iteration i

and

every lane in the tile

A versioning condition does not need to be mathematically necessary.

It may be conservative.

But it must be sufficient.

Condition false
→ optimization may be missed
→ original masked loop remains correct
Condition true without implying the mask
→ mask removed when some lanes are invalid
→ semantics change

The second outcome is an unsound compiler transformation.


The canonical tail mask has a recognizable form

CanonicalMaskValidator recognizes a loop-tail mask equivalent to:

range(0, END)
<
N - i * END

This is another way to write:

i * END + range(0, END)
<
N

The loop upper bound is expected to have the matching canonical form:

ceil(N / END)

represented as:

(N + END - 1) / END

The validator extracts two facts from the mask:

N
→ logical problem extent

END
→ tile width

It stores them as:

MaskInfo {
    N,
    END
}

The merged implementation recognizes the canonical loop mask only for the strict signed < form and checks that the loop upper bound uses the same logical values.


A canonical shape is only a candidate proof

Suppose the compiler sees:

loop upper bound:

(N_A + 31) / 32

and inside the loop:

mask:

range(0, 64)
<
N_B - i * 64

Both expressions resemble the canonical scheme.

upper bound
→ ceil-like division
mask
→ range compared with remaining elements

But they describe different partitions.

Loop:
N_A
END = 32
Mask:
N_B
END = 64

A scalar condition derived from:

N_A and 32

does not prove a mask defined by:

N_B and 64

The syntax belongs to the same family.

The symbols and tile width do not.


A numerical counterexample makes the mismatch visible

Take:

N_A = 64
END_A = 32

The loop upper bound is:

ceil(64 / 32)
=
2

A full-tile condition derived from that loop is true.

64 % 32
=
0

Now let another load in the same loop use:

N_B = 64
END_B = 64

Its mask is:

range(0, 64)
<
64 - i * 64

At the first iteration:

i = 0

range(0, 64) < 64
→ every lane true

At the second iteration:

i = 1

range(0, 64) < 0
→ every lane false

The loop-derived condition is still true:

N_A is divisible by 32

But the second mask is not true at all.

Loop guard
→ true

Mask B in iteration 1
→ false

Removing Mask B would execute lanes that the original program explicitly disabled.

Whether those accesses become physically out of bounds depends on the surrounding allocation and pointer calculation.

The semantic mismatch already exists before that question.


The regression contains exactly this symbolic separation

The new two-canonical-masks.mlir regression builds one loop whose upper bound is:

(%arg2 + 31) / 32

The first load uses:

N = %arg2
END = 32

Its mask corresponds to the loop upper bound.

The second load uses:

N = %arg3
END = 64

Its mask is unrelated to the condition generated from %arg2.

The test requires:

no versioning `scf.if`

both loads remain masked

It also contains a single-load case in which the mask alone uses %arg3 and 64, while the loop upper bound still uses %arg2 and 32.

Recognizing the arithmetic shape is deliberately not accepted as proof.


The old check matched structure without matching identity

Before the patch, hasCanonicalUpperBound() checked whether the loop upper bound looked like:

(X + C) / D

with:

D = C + 1

That was enough to recognize:

(X + END - 1) / END

as a ceil-division shape.

But it did not require:

X
=
the mask’s N

or:

D
=
the mask’s END

The pass therefore established:

This loop upper bound is canonical for some N and END.

It then used that as though it had established:

This loop upper bound is canonical
for the exact N and END used by this mask.

The second statement is stronger.

The old code had not proven it.


The fixed check compares both N and END

The merged implementation now requires:

final(mask.N)
==
final(loop_ub.N)

and:

mask.END
==
loop_ub.END

Only then can a condition generated from the loop upper bound be used to remove the mask.

Same `N`

+

same `END`

+

canonical arithmetic structure

→ versioning condition can describe the mask

The source comment states the reason directly:

The versioning condition is derived from the loop upper bound’s N and END.

It implies the mask only when the mask uses the same N and END.

SSA identity matters even when names look equivalent

At the source-language level, developers may describe two values as:

N

But compiler IR works with exact SSA values.

%arg2

%arg3

are different facts unless the compiler has a proof connecting them.

The pass cannot assume:

Both values represent tensor length.

Therefore they are equal.

Nor can it assume:

Both denominators are tile-like constants.

Therefore 32 and 64 are interchangeable.

A proof generated from one SSA value does not automatically transfer to another value with a similar role or name.

Same conceptual category
≠
same compiler fact

One scalar guard can protect several masks only when they share the proof

A loop can contain multiple masked loads.

Using one versioning branch for all of them is valid when every mask is derived from the same:

N

END

loop induction variable

predicate contract

Then the generated conditions are equivalent.

Guard derived from N and END
        ↓
implies Mask A
        ↓
implies Mask B
        ↓
all selected loads can become unmasked

If even one mask depends on another extent or tile width:

Guard
→ does not imply every collected mask

The pass must keep the original loop.

This is not an optimization failure.

It is the correct fail-closed outcome.


A constant loop upper bound did not prove a dynamic N

The second boundary involved a crash rather than an unsound mask drop.

The loop had a constant upper bound:

0 to 8

The load mask still depended on a dynamic kernel argument:

%arg1

Conceptually:

range(0, 32)
<
N_dynamic - i * 32

The old hasCanonicalUpperBound() saw:

loop UB is a constant

and attempted to determine whether it could be the folded result of:

(N + END - 1) / END

To do that, it unconditionally cast the defining operation of N to:

arith::ConstantIntOp

But a function argument has no defining operation.

N is a block argument

→ getDefiningOp() returns null

→ cast<ConstantIntOp>(null)

→ compiler abort

PR #7791 lists this as the second confirmed crash on valid MLIR.


The crash also exposed a missing proof

Avoiding the null cast was necessary.

It was not the whole issue.

A constant upper bound does not reveal which dynamic N produced it.

Take:

END = 32
loop upper bound = 8

For positive N, the expression:

ceil(N / 32)
=
8

is true for many values:

N = 225

N = 226

...

N = 255

N = 256

But those values do not have the same tail behavior.

N = 256

256 % 32
=
0

all eight tiles are complete

N = 255

255 % 32
=
31

the eighth tile has one invalid lane

Both produce:

loop upper bound = 8

The upper bound alone cannot distinguish them.

Same number of iterations

≠

same all-lanes-valid proof

A folded constant loses the information needed for versioning

Suppose an earlier compiler stage evaluated:

(N + 31) / 32

to:

8

If N was also a known constant, the pass could verify whether the folded value matched the canonical expression.

N known

END known

UB known

→ relationship can be checked

But when N remains dynamic:

UB known

N unknown

the pass cannot reconstruct the original divisibility state.

Many possible N values map to the same quotient.

The transformation must reject the optimization rather than guess which one was intended.


The fixed implementation rejects the dynamic case

The merged code now asks:

Does maskInfo.N come from a ConstantIntOp?

If the answer is no:

return false

The loop remains unversioned.

Dynamic N

+

constant loop UB

→ no proof that the UB is the folded canonical form for this N

→ keep the masked loop

The separate getVersioningCond() path asserts that N is constant only after hasCanonicalUpperBound() has guaranteed that precondition.

The code no longer treats a missing defining operation as a constant node.


The regression protects rejection, not a replacement condition

The canonical-mask-dynamic-n-const-ub.mlir test uses:

constant loop upper bound:
8

dynamic N:
kernel argument

tile width:
32

The expected transformed IR contains:

the original `scf.for`

the original masked `tt.load`

no versioning `scf.if`

The test does not ask the pass to invent a more complicated runtime proof.

It establishes a narrower contract:

When a constant upper bound cannot be related to the mask’s dynamic N, do not version the loop.


A boundary-shaped comparison was not necessarily a boundary proof

The third defect belonged to loop-invariant masks.

The pass recognized patterns shaped like:

splat(offset)
+
make_range(start, end)
cmp
uniform constant

For example:

offset + range(0, 32)
<
128

This is a common vector boundary check.

For ordered predicates, the complete vector condition can often be reduced to one endpoint.


Ordered predicates can use the maximum or minimum lane

Consider:

offset + range(0, 32)
<
bound

The vector values are:

offset + 0

offset + 1

...

offset + 31

Every lane is less than bound exactly when the maximum lane is less than bound.

offset + 31 < bound

is therefore a sufficient scalar condition for:

all lanes satisfy <

For a lower-bound comparison:

offset + range(0, 32)
>=
bound

the minimum lane determines the result.

offset + 0 >= bound

implies that every later lane is also greater than or equal to the bound.

The merged implementation’s boundary-condition builder uses:

maximum lane for `<` and `<=`

minimum lane for `>` and `>=`

after requiring a supported ordered predicate.


Equality does not have the same monotonic property

Now consider:

offset + range(0, 32)
==
128

Take:

offset = 97

The lanes are:

97, 98, 99, ..., 128

Only the final lane satisfies equality.

lane 31
→ true

lanes 0–30
→ false

A scalar condition such as:

offset + 31 == 128

does not prove:

every lane equals 128

It proves only:

the maximum lane equals 128

For a range containing distinct values, all lanes cannot equal one scalar bound.

Removing the vector mask under such a condition would turn:

one enabled lane

into:

all lanes enabled

ne also cannot use a single endpoint rule

Consider:

offset + range(0, 32)
!=
128

Take:

offset = 98

The lanes are:

98, 99, ..., 128, 129

The maximum lane is:

129

and:

129 != 128

is true.

But one interior lane is exactly:

128

Therefore:

31 lanes
→ true

1 lane
→ false

Checking one endpoint does not prove that the bound is absent from the full interval.

A correct scalar proof for ne would need to establish that the bound lies entirely outside the range.

Conceptually:

maximum lane < bound

or

minimum lane > bound

That is a two-sided disjunction.

The existing boundary-condition generator was not designed to construct that proof.

The safe fix was to reject ne.


Shape recognition and predicate semantics are separate checks

The expressions:

offset + range == bound
offset + range != bound
offset + range < bound

all share the same structural skeleton.

splat

+

make_range

comparison

uniform constant

But their all-lanes semantics differ.

Structure
→ identifies a candidate family
Predicate
→ determines whether one endpoint can prove the vector condition

The old isBoundaryCheckPattern() accepted the structure without first rejecting eq and ne.

The fixed implementation checks:

isSupportedBoundPredicate(predicate)

before accepting the pattern.

The source comment explicitly states that equality and inequality cannot be reduced to the scalar condition generated by the pass.


The eq and ne regressions retain the masks

The new eq-boundary-check-mask.mlir file contains two functions.

Equality case

(offset + range(0, 32)) == 128

Inequality case

(offset + range(0, 32)) != 128

For both, the test requires:

no versioning `scf.if`

original `scf.for` remains

masked `tt.load` remains

The test does not ask the compiler to implement a stronger equality or interval-exclusion proof.

It ensures that the existing endpoint-based proof is not applied outside its supported predicate family.


Rejecting the optimization is a correctness result

A compiler optimization can have three broad outcomes.

Proof established
→ transform
Proof unavailable
→ leave IR unchanged
Proof unavailable but transform anyway
→ unsound

The second outcome may produce a slower kernel.

It remains correct.

For #7791, the fixed pass deliberately rejects:

constant UB with dynamic N

mismatched mask and UB symbols

mismatched tile widths

eq/ne boundary predicates

The patch does not need to solve every possible proof problem.

It must stop claiming a proof where none exists.


“Canonical” does not mean “universally removable”

The word canonical can be misleading.

A canonical shape means:

the expression follows one recognized representation

It does not mean:

the transformation is always safe

The compiler still needs to verify:

  • the relevant predicate,

  • the exact SSA value used as N,

  • the exact tile width used as END,

  • the loop upper-bound relationship,

  • and every operation whose mask will be removed.

Canonical syntax

+
matching identities

+
supported predicate

+
implication proof

→ removable mask

Leaving out any term weakens the conclusion.


The proof must follow the value dependency, not the variable name

A compiler pass cannot rely on labels such as:

N

length

size

end

Two values may both represent dimensions in the source model and still differ at runtime.

The proof must follow:

which SSA value produced the loop bound

which SSA value produced the mask

which constants define their tile widths

which operations connect them

That is why the merged check uses normalized final values rather than matching only operation classes.


A versioning condition is a theorem about one expression

Suppose the pass creates:

N_A % 32 == 0

That condition is a theorem about:

N_A

and

32

It says nothing about:

N_B

64

another comparison predicate

another tensor’s boundary

The presence of all those expressions inside one loop does not merge their meanings.

Same loop body

≠

same proof domain

A scalar guard can remove several masks only when each mask is shown to be a consequence of that exact guard.


The pass now treats the fast path as a proven region

Loop versioning creates two semantic regions.

Then region

→ stronger assumptions hold

→ masks may be removed
Else region

→ assumptions do not hold

→ original masks remain

The then region is valid only if its branch condition establishes every assumption consumed inside it.

That makes the branch condition similar to a contract:

Enter this region only when
all selected vector accesses are fully valid.

Pattern matching finds candidate assumptions.

Validation proves them.


The tests operate entirely at the MLIR level

These regressions run:

triton-opt

with

-triton-intel-remove-masks

and inspect the output with FileCheck.

They require no Intel GPU.

They do not measure kernel latency or generated device instructions.

They directly test:

  • whether the compiler crashes,

  • whether an scf.if is created,

  • and whether masked loads remain masked.

The issue and PR describe all six failures as reproduced on valid MLIR input.


The proof regressions protect negative outcomes

Optimization tests often focus on successful transformations.

Input contains mask

Output no longer contains mask

The important tests here assert the opposite.

Input resembles optimization candidate

but proof is insufficient

→ output must keep the mask

Examples:

dynamic N with constant UB
→ no versioning
different N or END
→ no versioning
eq or ne
→ no versioning

These are negative regression contracts.

They protect the optimizer’s refusal boundary.


A negative test prevents future overgeneralization

A later developer may simplify code and observe:

Both masks have the same arithmetic shape.

Why not accept both?

The regression records the counterexample.

Shape equality
→ not enough

Likewise, someone may broaden supported comparison predicates:

`cmp` is already present.

Why not allow `eq` and `ne`?

The tests preserve the reason those predicates were excluded.

Endpoint condition
→ does not imply every lane

Negative tests preserve the knowledge that an apparently broader optimization is not necessarily sound.


The patch fixed proof construction, not runtime masking cost

PR #7791 does not publish a benchmark comparing:

masked loop

versus

versioned unmasked loop

It does not quantify how often the newly rejected cases occur in real kernels.

Its direct purpose was compiler correctness.

Crash path
→ reject rather than cast null
Unproven path
→ retain mask rather than remove it

Any performance impact from fewer versioned loops remains outside the published evidence.


The patch does not claim the rejected masks can never be optimized

An eq or ne mask could theoretically be optimized using a different proof or rewrite.

For example, an ne range can be proven all true when:

bound < minimum lane

or

bound > maximum lane

A compiler could build that disjunction.

PR #7791 does not do so.

The accurate statement is:

The existing scalar boundary-condition generator
does not prove eq/ne masks removable.

Not:

Equality masks can never be optimized.

Likewise, a dynamic N with a constant loop upper bound may become provable if additional range facts are available.

This pass did not possess those facts.

It therefore left the mask intact.


What the patch directly changed for the proof boundary

The merged changes directly:

  • reject constant loop upper bounds when the mask’s N is not a constant,

  • require the mask’s N to match the N used by the loop upper bound,

  • require the mask’s END to match the loop upper bound’s divisor,

  • reject unsupported eq and ne boundary predicates,

  • and add MLIR regressions requiring the original masked loops to remain.

The same PR also adds a mutation-free canVersion() precheck so all collected masks are validated before versioning begins.

That transformation-order issue is the subject of Part 3.


What this article does not establish

The public evidence does not establish that:

  • these unsound rewrites produced a published numerical failure on an Intel XPU,

  • every canonical mask is now exhaustively validated,

  • every ordered integer predicate is accepted by every validator,

  • all dynamic dimensions must always prevent versioning,

  • eq and ne can never be optimized by another proof,

  • retaining masks has no performance cost,

  • every RemoveMasks path is covered by these three regression files,

  • or the pass is formally verified.

The supported conclusion is narrower:

The old pass could treat a canonical-looking shape as sufficient evidence even when the generated guard referred to different symbolic bounds or an unsupported comparison predicate. The merged fix rejects those cases and preserves the original masks.


The complete proof failure chain

The unsound versioning path can be written as:

Mask resembles a recognized boundary form
        ↓
loop upper bound resembles canonical ceil division
        ↓
exact N / END identity not checked
or predicate semantics not supported
        ↓
scalar condition generated
        ↓
condition does not imply every vector lane
        ↓
mask removed in fast path
        ↓
original program semantics no longer guaranteed

The repaired path is:

Recognize candidate structure
        ↓
verify supported predicate
        ↓
extract mask N and END
        ↓
verify loop-UB N and END are identical
        ↓
reject dynamic folded cases that cannot be reconstructed
        ↓
only then generate versioning condition

The final lesson is that a proof has operands too

Compiler engineers often inspect an expression and recognize a familiar pattern.

ceil division

boundary mask

tile range

comparison with an extent

Recognition is useful.

It is not the final proof.

A proof also depends on:

which value is N

which constant is END

which predicate is used

which lanes the scalar condition represents

which operations will consume the assumption

Two expressions can have the same shape and still describe different facts.

Intel Triton XPU #7791 repaired RemoveMasks by requiring the versioning condition to belong to the exact mask it was used to remove.


Previous article

Intel Triton XPU #7791 — Why an Always-False Masked Load Without other Crashed RemoveMasks

Link the title above to Part 1 after publication.

Related material


Patch status: Merged into Intel Triton XPU main
Affected pass: TritonIntelRemoveMasks
Optimization: Loop versioning with mask removal in the guarded fast path
Canonical mask: range(0, END) < N - i * END
Required identity: The mask and loop upper bound must use the same N and END
Dynamic-constant boundary: A constant loop UB cannot be matched to a dynamic mask N without additional proof
Rejected predicates: eq and ne in the endpoint-based boundary-check path
Failure forms: One compiler crash and two unsound mask-removal classes
Regression form: Valid MLIR through triton-opt and FileCheck
Hardware requirement: None
Direct result: Unproven loops remain masked rather than being versioned
Evidence boundary: No XPU runtime-performance or production-model result is claimed

This is Part 2 of a three-part series on Intel Triton XPU’s RemoveMasks correctness boundary.

Part 1 examined why removing an always-false load still required a valid typed SSA replacement when its optional other operand was absent.

This article examined why canonical arithmetic shape did not prove that a loop condition implied a mask, and why dynamic dimensions, mismatched symbolic bounds, and equality predicates had to be rejected.

Part 3 examines why every collected mask had to be validated before the first IR mutation, why a result-less scf.if could not receive a second explicit terminator, and why versioning an operation while walking the same IR required an explicit walker-lifetime boundary.

#Intel #Triton #XPU #MLIR #CompilerOptimization #LoopVersioning #MaskedLoad #RangeAnalysis #ProofObligation #CompilerCorrectness #CodeAnalysis

Popular posts from this blog

AMD Is Trying to Turn Kernel Diversity Into One Software Asset

PyTorch #188031 — Why Did Forward Use 64-Bit Indexing While Backward Still Used uint32_t?

NVIDIA CUTLASS #3017 — Why load() and store() Mapped the Same Tile Coordinate to Different Addresses