Skip to content

Infer src/dst of allowReorder reshapes#9926

Merged
neildhar merged 2 commits into
mainfrom
neildhar/pr9926
Apr 7, 2026
Merged

Infer src/dst of allowReorder reshapes#9926
neildhar merged 2 commits into
mainfrom
neildhar/pr9926

Conversation

@neildhar
Copy link
Copy Markdown
Collaborator

@neildhar neildhar commented Apr 4, 2026

Always infer the src/dst of reshapes, even if allowReorder is set. The
result is valid for allowReorder reshapes, even if there isn't a single
canonical encoding. When the existing encoding is one of the possible
results, we prefer that to minimize changes.

This allows inference to always succeed on reshapes, and any heuristics
on whether to use the inferred value can be maintained by the caller.

One example I identified while looking at this was that allowReorder
reshapes will currently fail backward remat in RemoveLayoutConversions
if the reshape cannot be rematerialised with the same source encoding.
This PR instead changes RemoveLayoutConversions to check specifically
for whether the reshape has been marked as efficient, and otherwise
just do the remat. (this is a potentially perf sensitive change)


@neildhar neildhar changed the base branch from main to neildhar/pr9924 April 4, 2026 03:06
@neildhar
Copy link
Copy Markdown
Collaborator Author

neildhar commented Apr 4, 2026

This is intended to supersede #9906 which somehow got merged while trying to update my stack. See the discussion there.

neildhar added 2 commits April 4, 2026 16:11
`inferDstEncoding` currently tries to unconditionally build a sliced
encoding from the src encoding it is given. But this is incorrect if the
source is rank 1, since we can't take a slice of a rank 1 tensor.
Always infer the src/dst of reshapes, even if allowReorder is set. The
result is valid for allowReorder reshapes, even if there isn't a single
canonical encoding. When the existing encoding is one of the possible
results, we prefer that to minimize changes.

This allows inference to always succeed on reshapes, and any heuristics
on whether to use the inferred value can be maintained by the caller.

One example I identified while looking at this was that allowReorder
reshapes will currently fail backward remat in RemoveLayoutConversions
if the reshape cannot be rematerialised with the same source encoding.
This PR instead changes RemoveLayoutConversions to check specifically
for whether the reshape has been marked as efficient, and otherwise
just do the remat. (this is a potentially perf sensitive change)
Base automatically changed from neildhar/pr9924 to main April 6, 2026 22:04
@neildhar neildhar marked this pull request as ready for review April 6, 2026 22:05
@neildhar neildhar requested review from ThomasRaoux and lezcano April 6, 2026 22:05
@neildhar neildhar merged commit 5f96878 into main Apr 7, 2026
20 of 23 checks passed
@neildhar neildhar deleted the neildhar/pr9926 branch April 7, 2026 00:38
ThomasRaoux added a commit to ThomasRaoux/triton that referenced this pull request Apr 9, 2026
ThomasRaoux added a commit that referenced this pull request Apr 9, 2026
neildhar added a commit that referenced this pull request Apr 10, 2026
Always infer the src/dst of reshapes, even if allowReorder is set. The
result is valid for allowReorder reshapes, even if there isn't a single
canonical encoding. When the existing encoding is one of the possible
results, we prefer that to minimize changes.

This allows inference to always succeed on reshapes, and any heuristics
on whether to use the inferred value can be maintained by the caller.

One example I identified while looking at this was that allowReorder
reshapes will currently fail backward remat in RemoveLayoutConversions
if the reshape cannot be rematerialised with the same source encoding.
This PR instead changes RemoveLayoutConversions to check specifically
for whether the reshape has been marked as efficient, and otherwise
just do the remat. (this is a potentially perf sensitive change)
neildhar added a commit that referenced this pull request Apr 10, 2026
Reland of #9926.

Always infer the src/dst of reshapes, even if allowReorder is set. The
result is valid for allowReorder reshapes, even if there isn't a single
canonical encoding. When the existing encoding is one of the possible
results, we prefer that to minimize changes.

This allows inference to always succeed on reshapes, and any heuristics
on whether to use the inferred value can be maintained by the caller.

One example I identified while looking at this was that allowReorder
reshapes will currently fail backward remat in RemoveLayoutConversions
if the reshape cannot be rematerialised with the same source encoding.
This PR instead changes RemoveLayoutConversions to check specifically
for whether the reshape has been marked as efficient, and otherwise
just do the remat. (this is a potentially perf sensitive change)
plognjen pushed a commit to plognjen/triton that referenced this pull request Apr 14, 2026
Always infer the src/dst of reshapes, even if allowReorder is set. The
result is valid for allowReorder reshapes, even if there isn't a single
canonical encoding. When the existing encoding is one of the possible
results, we prefer that to minimize changes.

This allows inference to always succeed on reshapes, and any heuristics
on whether to use the inferred value can be maintained by the caller.

One example I identified while looking at this was that allowReorder
reshapes will currently fail backward remat in RemoveLayoutConversions
if the reshape cannot be rematerialised with the same source encoding.
This PR instead changes RemoveLayoutConversions to check specifically
for whether the reshape has been marked as efficient, and otherwise
just do the remat. (this is a potentially perf sensitive change)
plognjen pushed a commit to plognjen/triton that referenced this pull request Apr 14, 2026
neildhar added a commit that referenced this pull request Apr 14, 2026
Reland of #9926.

Always infer the src/dst of reshapes, even if allowReorder is set. The
result is valid for allowReorder reshapes, even if there isn't a single
canonical encoding. When the existing encoding is one of the possible
results, we prefer that to minimize changes.

This allows inference to always succeed on reshapes, and any heuristics
on whether to use the inferred value can be maintained by the caller.

One example I identified while looking at this was that allowReorder
reshapes will currently fail backward remat in RemoveLayoutConversions
if the reshape cannot be rematerialised with the same source encoding.
This PR instead changes RemoveLayoutConversions to check specifically
for whether the reshape has been marked as efficient, and otherwise
just do the remat. (this is a potentially perf sensitive change)
neildhar added a commit that referenced this pull request Apr 15, 2026
Reland of #9926.

Always infer the src/dst of reshapes, even if allowReorder is set. The
result is valid for allowReorder reshapes, even if there isn't a single
canonical encoding. When the existing encoding is one of the possible
results, we prefer that to minimize changes.

This allows inference to always succeed on reshapes, and any heuristics
on whether to use the inferred value can be maintained by the caller.

One example I identified while looking at this was that allowReorder
reshapes will currently fail backward remat in RemoveLayoutConversions
if the reshape cannot be rematerialised with the same source encoding.
This PR instead changes RemoveLayoutConversions to check specifically
for whether the reshape has been marked as efficient, and otherwise
just do the remat. (this is a potentially perf sensitive change)
raymondtay pushed a commit to raymondtay/triton that referenced this pull request Apr 18, 2026
Reland of triton-lang#9926.

Always infer the src/dst of reshapes, even if allowReorder is set. The
result is valid for allowReorder reshapes, even if there isn't a single
canonical encoding. When the existing encoding is one of the possible
results, we prefer that to minimize changes.

This allows inference to always succeed on reshapes, and any heuristics
on whether to use the inferred value can be maintained by the caller.

One example I identified while looking at this was that allowReorder
reshapes will currently fail backward remat in RemoveLayoutConversions
if the reshape cannot be rematerialised with the same source encoding.
This PR instead changes RemoveLayoutConversions to check specifically
for whether the reshape has been marked as efficient, and otherwise
just do the remat. (this is a potentially perf sensitive change)
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