Skip LLVM jump threading for functions with convergent ops

[39ali]

Fixes #11148.

What happens

The kernel in #11148 crashes with an illegal memory access on sm_86 and
silently returns wrong results on sm_120.

Why

Some GPU instructions (bar.sync, ldmatrix) need all 32 threads of a warp
to arrive together. When a branch temporarily splits a warp, ptxas re-joins
the threads at the point where the branch's paths merge back — and it finds
those merge points by looking at the shape of the control-flow graph.

LLVM's jump-threading pass optimizes branches by rerouting paths around
redundant condition tests. Doing that to this kernel destroys the merge
point in front of a bar.sync/ldmatrix pair: after the rewrite, ptxas has
no valid place left to re-join the warp, so the ldmatrix runs with a
half-joined warp. That's undefined behavior.

Jump threading already refuses to duplicate blocks that contain these
warp-synchronous ops, but nothing stops it from rerouting edges into them.
That gap is the bug.

The fix

Don't run jump threading on functions that contain warp-synchronous
(convergent) ops. Functions without them — e.g. elementwise kernels — keep
it. Benchmarks (matmul, softmax, elementwise) show no measurable perf
change: GPU kernels have little of the branchy code this pass optimizes.

The proper fix is teaching jump threading itself to skip only the illegal
rewrites; I'm filing that upstream with a standalone repro. This skip can be
removed once Triton's LLVM includes it.

Performance

Benchmarked on RTX 5070 Ti (matmul 2048³ fp16, softmax 8192×2048, 16M
elementwise): all deltas within run-to-run noise (±2%, with the
identical-binary control kernel showing the same spread). Threading's wins
come from branchy scalar code that performance-sensitive GPU kernels avoid;
the branch-rich code it does find here (inlined libdevice bodies) is exactly
the material that miscompiles.

New contributor declaration

• I am not making a trivial change, such as fixing a typo in a comment.

• I have written a PR description following these
  rules.

• I have run pre-commit run --from-ref origin/main --to-ref HEAD.

• Select one of the following.

  • I have added tests.
    • /test for lit tests
    • /unittest for C++ tests
    • /python/test for end-to-end tests
  • This PR does not need a test because FILL THIS IN.

• Select one of the following.

  • I have not added any lit tests.
  • The lit tests I have added follow these best practices,
    including the "tests should be minimal" section. (Usually running Python code
    and using the instructions it generates is not minimal.)

[ThomasRaoux]

doesn't LLVM already handle that? The pass should not be able to apply jump threading when some ops are convergent?

[39ali]

@ThomasRaoux That's what I expected too, but the pass's convergent check is narrower than
that: it only refuses when the block it wants to duplicate contains
convergent ops. It does not refuse to thread just because convergent ops exist
nearby.

In this kernel the block it duplicates is a small guard block with no
convergent ops in it, so the check passes. The convergent block (the
bar.sync/ldmatrix from the layout conversion) is never duplicated it,
it just ends up with a new incoming edge from another region of the CFG. Nothing
in the pass validates that, but the extra edge removes the merge point ptxas
needs to reconverge the warp before the ldmatrix, which then executes
partially converged.

Verified by bisection: -opt-bisect-limit flips the breakage exactly at
jump-threading, and plain opt -passes='default<O3>' on the pre-optimization
IR reproduces it outside Triton. I'll file the missing check upstream; this PR
is the interim mitigation until the LLVM pin includes it.

[ThomasRaoux]

@ThomasRaoux That's what I expected too, but the pass's convergent check is narrower than that: it only refuses when the block it wants to duplicate contains convergent ops. It does not refuse to thread just because convergent ops exist nearby.

In this kernel the block it duplicates is a small guard block with no convergent ops in it, so the check passes. The convergent block (the bar.sync/ldmatrix from the layout conversion) is never duplicated it, it just ends up with a new incoming edge from another region of the CFG. Nothing in the pass validates that, but the extra edge removes the merge point ptxas needs to reconverge the warp before the ldmatrix, which then executes partially converged.

Verified by bisection: -opt-bisect-limit flips the breakage exactly at jump-threading, and plain opt -passes='default<O3>' on the pre-optimization IR reproduces it outside Triton. I'll file the missing check upstream; this PR is the interim mitigation until the LLVM pin includes it.

why would the extra edge remove the sync point?

Sounds like fixing it in LLVM is the way to go.

[39ali]

@ThomasRaoux it's a triton issue, target triple should be set for jump threading to work properly, i updated the pr and pushed a new fix

see : 213870

[neildhar]

I suspect this bug may be ptxas, I was able to minimize this reproducer such that it produces incorrect results even without jump threading. (instead depending purely on ptxas optimization level)

That said, it seems like setting the target triple is the right thing to do either way. But since it can change the final generated code substantially, it may be disruptive.

[neildhar]

I think we can should add a separate attach_target_triple here in the shared code and remove the AMD specific one we currently have.

[ThomasRaoux]

this would be the right fix however I believe this will impact the vectorizer and cause performance regressions

[antoniofrighetto]

This is indeed a bit unfortunate, though I'd be inclined to think that NVPTX TTI returning suboptimal costs is an orthogonal issue that might be worth investigating separately. It could be helpful to understand whether the cost model could be tuned to produce more accurate vectorization decisions.

[39ali]

what's the move here ? should i implement the attach_target_triple refactor and then we should investigate performance regressions separately ?

[neildhar]

Let's do the refactor first, I think we should be good but we may have to revert if we hit issues.

[39ali]

should be good to go

[neildhar]

LGTM, could you update the title and summary to reflect what the change actually does now?

[neildhar]

I don't think this diagnosis is correct, because disabling ptxas optimizations also fixes the bug. The bug is likely to actually be in ptxas.

[39ali]

@neildhar should i keep this test even ? and can you share the code that would reproduce the bug without jump threading for future reference ?

[neildhar]

We can drop it.

Sure, here is the repro:

import torch
from triton.experimental import gluon as g
from triton.experimental.gluon import language as gl


@g.jit
def kernel(inputs, outputs):
    lo = gl.load(inputs + 1)
    cond_false = gl.load(inputs) != 0
    src = gl.arange(0, 4, layout=gl.BlockedLayout([4], [32], [1], [0]))
    dst = gl.arange(0, 4, layout=gl.BlockedLayout([1], [32], [1], [0]))
    if lo != 0:
        if cond_false:
            gl.store(outputs + 4, gl.sin(lo))
            return
        gl.store(outputs + 4, gl.sin(lo))
    elif cond_false:
        gl.store(outputs + 4, gl.sin(lo))
        return
    value = gl.load(inputs + src)
    value = gl.convert_layout(value, gl.BlockedLayout([1], [32], [1], [0]))
    gl.store(outputs + 3 - dst, value)


inputs = torch.arange(4, device="cuda", dtype=torch.float32)
outputs = torch.empty(5, device="cuda", dtype=torch.float32)
compiled = kernel[(1,)](inputs, outputs, num_warps=1)
torch.cuda.synchronize()
matches = torch.equal(outputs[:4], inputs.flip(0))
print(
    f"matches={matches} shared={compiled.metadata.shared} "
    f"ldmatrix={compiled.asm['ptx'].count('ldmatrix')} "
    f"ttgir_bytes={len(compiled.asm['ttgir'])}"
)
assert matches, outputs