fix(mlir): loosen restrictions on folding dynamic reshapes #8
Conversation
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
AGindinson
force-pushed
the
artem/expand-of-collapse
branch
2 times, most recently
from
b7f9d88 to
21ff423
Compare
This stack of pull requests is managed by Graphite. Learn more about stacking. |
The main idea behind the change is to allow expand-of-collapse folds for reshapes like `?x?xk` -> `?` (k>1). The rationale here is that the expand op must have a coherent index/affine expression specified in its `output_shape` argument (see example below), and if it doesn't, the IR has already been invalidated at an earlier stage: ``` %c32 = arith.constant 32 : index %div = arith.divsi %<some_index>, %c32 : index %collapsed = tensor.collapse_shape %41#1 [[0], [1, 2], [3, 4]] : tensor<9x?x32x?x32xf32> into tensor<9x?x?xf32> %affine = affine.apply affine_map<()[s0] -> (s0 * 32)> ()[%div] %expanded = tensor.expand_shape %collapsed [[0], [1, 2], [3]] output_shape [9, %div, 32, %affine] : tensor<9x?x?xf32> into tensor<9x?x32x?xf32> ``` On the above assumption, adjust the routine in `getReassociationIndicesForCollapse()` to allow dynamic reshapes beyond just `?x..?x1x1x..x1` -> `?`. Moreover, the reassociation util was refactored to clearly distinguish between dynamic and static subshapes. A few known caveats were noted as a comment; I don't think it's possible to fold all qualifying dynamic shape patterns in a deterministic way without looking into affine expressions simultaneously (which would be difficult to maintain in a single general utility for all reliant passes, and would therefore require a larger refactor).
AGindinson
force-pushed
the
artem/expand-of-collapse
branch
from
21ff423 to
6762ea2
Compare
chrsmcgrr
approved these changes
Apr 30, 2025
maxbartel
approved these changes
Apr 30, 2025
|
The placement of the file is super weird btw. I guess Tensor and Memref are using this and for historic reasons it is in this folder. Nowadays people would implement an interface I guess |
|
Yeah, all the cases that I've debugged through were pretty much Tensor-specific. Even if it's significant for Memref, I'd have an argument or two for making implementations dialect- or even pattern-specific, looking at how |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
The main idea behind the change is to allow expand-of-collapse folds for reshapes like
?x?xk->?(k>1). The rationale here is that the expand op must have a coherent index/affine expression specified in itsoutput_shapeargument (see example below), and if it doesn't, the IR has already been invalidated at an earlier stage:On the above assumption, adjust the routine in
getReassociationIndicesForCollapse()to allow dynamic reshapes beyond?x...x?x1x...x1->?.Moreover, the reassociation util was refactored to clearly distinguish between dynamic and static subshapes. A few known caveats were noted as a comment; I don't think it's possible to fold all qualifying dynamic shape patterns in a deterministic way without looking into affine expressions simultaneously (which would be difficult to maintain in a single general utility for all reliant passes, and would therefore require a larger refactor).