forked from apache/tvm
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[TIR][Schedule] Transform layout quality of life (apache#11269)
* [TIR][Schedule] Added Schedule.transform_layout_sugared * [TE][TIR] Reduced duplication in TE/TIR layout transformations Previously, the implementations of `tir.IndexMap.from_func` and `te.Stage.transform_layout` had significant duplication to handle argument parsing. This commit extracts the shared logic into `tir.IndexMap`. * Enabled *args in Schedule.transform_layout_sugared * Fix lint error * Allow Schedule.transform_layout_sugared to set axis separators * Merged transform_layout_sugared functionality into transform_layout * Fix lint errors * Fix lint error * Fixed docstring errors * Updated/tested TransformatLayoutTraits::UnpackedAsPython * Disabled exec-used check for running trace.as_python() * Updated SetAxisSeparatorTraits::UnpackedAsPython * Updated unit test that was added in merge commit * Fixed the argument name for TensorizeTraits This wasn't checked before, but was the only other issue caught by the updates to verify_trace_roundtrip. * Re-enable type checks of transform_layout/set_axis_separator Disabled while waiting for apache#11289, which was required for the `Tuple` argument. * Updated a few additional transform_layout usages from main
- Loading branch information
1 parent
e0a4e81
commit d8b1fd1
Showing
9 changed files
with
385 additions
and
132 deletions.
There are no files selected for viewing
This file contains 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 file contains 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
Oops, something went wrong.