[PCF] Add tensor.collapse_shape fusion into pcf.generic/loop#23491
Merged
qedawkins merged 4 commits intoiree-org:mainfrom Mar 5, 2026
Merged
[PCF] Add tensor.collapse_shape fusion into pcf.generic/loop#23491qedawkins merged 4 commits intoiree-org:mainfrom
qedawkins merged 4 commits intoiree-org:mainfrom
Conversation
qedawkins
force-pushed
the
fuse-collapse-shape
branch
from
59a716b to
506dcb1
Compare
Fuse tensor.collapse_shape consumers into producer pcf.generic/pcf.loop ops. This supports all possible dynamic cases by emitting loops along outer dimensions, effectively unrolling the producer `pcf.write_slice` ops allowing for continued fusion. The alternative to this (essentially premature unrolling) is to put the reshape on the sref, though to continue fusing we still would have to do the same (or propagate the collapse, which should have already been attempted by this point). Currently this pattern fails on multi-use results to avoid situations where we end up duplicating writes for a reshape. This is not a great solution but the easiest thing to do right now until a better option presents itself. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
qedawkins
force-pushed
the
fuse-collapse-shape
branch
from
506dcb1 to
8c78a66
Compare
krzysz00
reviewed
Feb 28, 2026
Max191
approved these changes
Mar 5, 2026
Contributor
Max191
left a comment
Max191
left a comment
There was a problem hiding this comment.
Overall this LGTM % some nits.
There are a lot of inline comments that probably aren't necessary here, though (I just pointed out a few, but there are a bunch more). Would be worth having Claude run through it and remove them.
I didn't carefully review the logic for correctness, but the test coverage is good, so I'm going to trust it. Maybe just add a fully dynamic collapse_shape test.
Please also make sure that you are using the affine.linearize/delinearize ops where possible.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fuse tensor.collapse_shape consumers into producer pcf.generic/pcf.loop ops. This supports all possible dynamic cases by emitting loops along outer dimensions, effectively unrolling the producer
pcf.write_sliceops allowing for continued fusion. The alternative to this (essentially premature unrolling) is to put the reshape on the sref, though to continue fusing we still would have to do the same (or propagate the collapse, which should have already been attempted by this point).