[TIR][Arith] Implemented padded inverses in IndexMap #11235
Merged
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
This allows it to be used for any expression containing an `IterMapExpr`, not just expressions whose top-level node is an `IterMapExpr`.
The existing DetectIterMap tries to rewrite index expression as a linear combination of split/fused iterators, where the new iterators cover the exact same indices as the original expression. DetectPaddedIterMap relaxes this condition, allowing the new iterators to cover a superset of indices that the initial index expression covered. It uses the minimum amount of padding necessary to represent these transformations, and also a predicate that identifies any padding that has been added. This is a utility function to be used for layout transformations of buffers, in cases where the pre-transformation shape of the buffer does not evenly fit into the post-transformation shape.
Allow non-surjective transformations, with DetectIterMap used to determine the minimum padding to insert. Returns the inverse function, along with a predicate that identifies padding indices. The predicate is in terms of the transformed variables.
- `IndexMap::Inverse` exposed as `IndexMap.inverse` - `IndexMap::MapShape` exposed as `IndexMap.map_shape` - `IndexMap::NonSurjectiveInverse` exposed as `IndexMap.non_surjective_inverse`
In preparation for adding additional tests for the IndexMap class, which will require this functionality.
Initially disabled as dynamic shapes resulted in padded lengths whose divisiblity couldn't be proven. Re-enabled along with a simplification rule to resolve it.
Contributor
Author
|
Also, name suggestions for "non-surjective inverse" would be appreciated. It's something of a mouthful, and describes what it is doing without describing why it would be desired. |
Member
|
Thanks for the PR, I like this change. There is a failing test case |
Contributor
Author
Thank you, and I'm liking it as well. This sort of utility should be useful regardless on the buffer padding semantics used.
Yeah, so far I've narrowed it down to the handling of cases like |
Contributor
Hi~ Does it mean pass index like |
Contributor
Author
Prior to this PR, that is correct. The check on the main branch is here, which requires the offset in a numerator to be zero. After this PR, an index like |
vinx13
approved these changes
May 13, 2022
mehrdadh
pushed a commit
to mehrdadh/tvm
that referenced
this pull request
May 16, 2022
* [Debug] Error logging in DetectIterMap * [Affine] Allowed PrimExpr argument to NormalizeIterMapToExpr This allows it to be used for any expression containing an `IterMapExpr`, not just expressions whose top-level node is an `IterMapExpr`. * [Affine] Implemented DetectPaddedIterMap The existing DetectIterMap tries to rewrite index expression as a linear combination of split/fused iterators, where the new iterators cover the exact same indices as the original expression. DetectPaddedIterMap relaxes this condition, allowing the new iterators to cover a superset of indices that the initial index expression covered. It uses the minimum amount of padding necessary to represent these transformations, and also a predicate that identifies any padding that has been added. This is a utility function to be used for layout transformations of buffers, in cases where the pre-transformation shape of the buffer does not evenly fit into the post-transformation shape. * [IndexMap] Implemented IndexMap::NonSurjectiveInverse Allow non-surjective transformations, with DetectIterMap used to determine the minimum padding to insert. Returns the inverse function, along with a predicate that identifies padding indices. The predicate is in terms of the transformed variables. * [IndexMap] Exposed methods to python - `IndexMap::Inverse` exposed as `IndexMap.inverse` - `IndexMap::MapShape` exposed as `IndexMap.map_shape` - `IndexMap::NonSurjectiveInverse` exposed as `IndexMap.non_surjective_inverse` * [IndexMap] Extracted _assert_equal_index_map into class method In preparation for adding additional tests for the IndexMap class, which will require this functionality. * [IndexMap] Added unit tests for new behavior * Re-enabled divisibility check in CheckMapping Initially disabled as dynamic shapes resulted in padded lengths whose divisiblity couldn't be proven. Re-enabled along with a simplification rule to resolve it. * Fixed breakage in compute_at primitive * Corrected typos/examples in docstring
shtinsa
pushed a commit
to Deelvin/tvm
that referenced
this pull request
May 17, 2022
* [Debug] Error logging in DetectIterMap * [Affine] Allowed PrimExpr argument to NormalizeIterMapToExpr This allows it to be used for any expression containing an `IterMapExpr`, not just expressions whose top-level node is an `IterMapExpr`. * [Affine] Implemented DetectPaddedIterMap The existing DetectIterMap tries to rewrite index expression as a linear combination of split/fused iterators, where the new iterators cover the exact same indices as the original expression. DetectPaddedIterMap relaxes this condition, allowing the new iterators to cover a superset of indices that the initial index expression covered. It uses the minimum amount of padding necessary to represent these transformations, and also a predicate that identifies any padding that has been added. This is a utility function to be used for layout transformations of buffers, in cases where the pre-transformation shape of the buffer does not evenly fit into the post-transformation shape. * [IndexMap] Implemented IndexMap::NonSurjectiveInverse Allow non-surjective transformations, with DetectIterMap used to determine the minimum padding to insert. Returns the inverse function, along with a predicate that identifies padding indices. The predicate is in terms of the transformed variables. * [IndexMap] Exposed methods to python - `IndexMap::Inverse` exposed as `IndexMap.inverse` - `IndexMap::MapShape` exposed as `IndexMap.map_shape` - `IndexMap::NonSurjectiveInverse` exposed as `IndexMap.non_surjective_inverse` * [IndexMap] Extracted _assert_equal_index_map into class method In preparation for adding additional tests for the IndexMap class, which will require this functionality. * [IndexMap] Added unit tests for new behavior * Re-enabled divisibility check in CheckMapping Initially disabled as dynamic shapes resulted in padded lengths whose divisiblity couldn't be proven. Re-enabled along with a simplification rule to resolve it. * Fixed breakage in compute_at primitive * Corrected typos/examples in docstring
shingjan
pushed a commit
to shingjan/tvm
that referenced
this pull request
May 17, 2022
* [Debug] Error logging in DetectIterMap * [Affine] Allowed PrimExpr argument to NormalizeIterMapToExpr This allows it to be used for any expression containing an `IterMapExpr`, not just expressions whose top-level node is an `IterMapExpr`. * [Affine] Implemented DetectPaddedIterMap The existing DetectIterMap tries to rewrite index expression as a linear combination of split/fused iterators, where the new iterators cover the exact same indices as the original expression. DetectPaddedIterMap relaxes this condition, allowing the new iterators to cover a superset of indices that the initial index expression covered. It uses the minimum amount of padding necessary to represent these transformations, and also a predicate that identifies any padding that has been added. This is a utility function to be used for layout transformations of buffers, in cases where the pre-transformation shape of the buffer does not evenly fit into the post-transformation shape. * [IndexMap] Implemented IndexMap::NonSurjectiveInverse Allow non-surjective transformations, with DetectIterMap used to determine the minimum padding to insert. Returns the inverse function, along with a predicate that identifies padding indices. The predicate is in terms of the transformed variables. * [IndexMap] Exposed methods to python - `IndexMap::Inverse` exposed as `IndexMap.inverse` - `IndexMap::MapShape` exposed as `IndexMap.map_shape` - `IndexMap::NonSurjectiveInverse` exposed as `IndexMap.non_surjective_inverse` * [IndexMap] Extracted _assert_equal_index_map into class method In preparation for adding additional tests for the IndexMap class, which will require this functionality. * [IndexMap] Added unit tests for new behavior * Re-enabled divisibility check in CheckMapping Initially disabled as dynamic shapes resulted in padded lengths whose divisiblity couldn't be proven. Re-enabled along with a simplification rule to resolve it. * Fixed breakage in compute_at primitive * Corrected typos/examples in docstring
Lunderberg
added a commit
to Lunderberg/tvm
that referenced
this pull request
May 23, 2022
Follow-up from apache#11235, all error messages should be based on expressions that are not IterMapExpr.
junrushao
pushed a commit
that referenced
this pull request
May 23, 2022
Follow-up from #11235, all error messages should be based on expressions that are not IterMapExpr.
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.
The goal of this PR is to allow
IndexMapto express transformations that introduce padding in the transformed shape. For an arbitrary input shape, the new methodIndexMap.non_surjective_inversedetermines the inverse transformation, along with a predicate specifying which coordinates in the transformed index space do not contain an inverse in the original index space. The previous behavior ofIndexMap.inverse, requiring transformations to be bijective over the range given, is maintained.This functionality will be used in the future to allow buffer transformations (see #9727 and #10538) to introduce padding to the buffer.