Skip to content

Handle region control flow in remat cost calculation#9201

Merged
neildhar merged 1 commit into
triton-lang:mainfrom
neildhar:single-use-yield
May 11, 2026
Merged

Handle region control flow in remat cost calculation#9201
neildhar merged 1 commit into
triton-lang:mainfrom
neildhar:single-use-yield

Conversation

@neildhar
Copy link
Copy Markdown
Collaborator

@neildhar neildhar commented Jan 13, 2026

Determining whether an op is used outside of the slice being
rematerialised is currently conservative. For instance, if an op is used
by an scf.yield, we automatically treat that as a non-slice user, even
if the value that the yield flows into is actually part of the slice or
completely unused.

To address this, trace through region control flow when deciding whether
an op is single use.

@neildhar neildhar force-pushed the single-use-yield branch 3 times, most recently from d365fbe to 88afe54 Compare January 19, 2026 03:21
peterbell10 pushed a commit that referenced this pull request Feb 16, 2026
…9245)

Replace the current recursive `isOpSingleUse` helper with an iterative
traversal that starts from the ops used directly outside a slice and
propagates to their operands. This simplifies the code and makes the
handling of block arguments explicit. This is an incremental step
towards properly handling block arguments (at least for for-loops) in
#9201.
@neildhar neildhar force-pushed the single-use-yield branch 6 times, most recently from ed21d2d to afc5e74 Compare February 22, 2026 23:14
@neildhar neildhar force-pushed the single-use-yield branch 4 times, most recently from 91d5f9d to 9a840ac Compare May 1, 2026 04:14
@neildhar neildhar changed the title [WIP] Improve detecting ops that will not survive after rematerialisation Handle region control flow in remat cost calculation May 1, 2026
@neildhar neildhar marked this pull request as ready for review May 1, 2026 04:20
@neildhar neildhar requested a review from ptillet as a code owner May 1, 2026 04:20
@neildhar neildhar force-pushed the single-use-yield branch from 9a840ac to efc5d78 Compare May 11, 2026 16:48
Determining whether an op is used outside of the slice being
rematerialised is currently conservative. For instance, if an op is used
by an `scf.yield`, we automatically treat that as a non-slice user, even
if the value that the yield flows into is actually part of the slice or
completely unused.

To address this, trace through region control flow when deciding whether
an op is single use.
@neildhar neildhar force-pushed the single-use-yield branch from efc5d78 to e2c4ef9 Compare May 11, 2026 17:55
Copy link
Copy Markdown
Collaborator

@ThomasRaoux ThomasRaoux left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@neildhar neildhar merged commit 25d2d2d into triton-lang:main May 11, 2026
9 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants