[TIR][Schedule] Transform layout quality of life by Lunderberg · Pull Request #11269 · apache/tvm

Lunderberg · 2022-05-10T21:33:18Z

Draft PR, seeing whether some of the TE-specific interface features can be pulled over into the TIR-based schedules. These are implemented as Schedule.transform_layout_sugared, which wraps around the existing Schedule.transform_layout.

Specify buffer by name, rather than by access_type/index.
Optional specify block by name, rather than needing a separate call to Schedule.get_block.
Enable *args in the index mapping function.
Allow AXIS_SEPARATOR in the index mapping function, delegate to Schedule.set_axis_separators

Lunderberg · 2022-05-10T21:35:15Z

@vinx13 A draft of some ease-of-use changes I'd been experimenting with this afternoon. These are python-only, and unpack to the same C++ calls, so it (hopefully) shouldn't cause issues elsewhere.

vinx13 · 2022-05-10T22:05:14Z

This is super cool, thanks for the PR. Shall we consider more informative names on what's being sugared. For example, axis_separators are related to physical layout, do you think something like transform_to_physical_layout would be more meaningful? Feel free to come up with better names

Lunderberg · 2022-05-11T14:38:49Z

I'd definitely be up for a better name, as transform_layout_sugared definitely leaves the reader uncertain what is being done. My one concern with having a completely different name is that it may not be clear to a reader that this is the same underlying schedule primitive in C++.

Another alternative would be to merge this as the input handling of Schedule.transform_layout. There's already some handling, such as the conversion from Callable to IndexMap, so this would be similar python-only handling. This would also avoid the issue of having multiple functions that aren't immediately obvious as wrappers around each other.

Lunderberg · 2022-05-11T17:05:03Z

Something sort of like the changes here are what I'm picturing. Not complete, because Union[Tuple[int,str],str] breaks the @type_checked annotation, but something sort of like this.

Lunderberg · 2022-05-12T14:37:01Z

#11289 should resolve the type-checking annotation issue with Tuples, if we want to go forward with the other changes here.

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`.

Lunderberg · 2022-05-16T13:49:28Z

CI functionality tests all passed, updates to resolve sphinx documentation tests, and rebase to resolve merge conflict.

vinx13 · 2022-05-17T02:59:16Z

- buffer_index: int,
- buffer_index_type: str,
+ block: Union[BlockRV, str],
+ buffer: Union[Tuple[str, int], str, Buffer],


this will break trace to python conversion. It is implemented here https://github.com/apache/tvm/blob/main/src/tir/schedule/primitive/layout_transformation.cc#L288. Unfortunately CI didn’t catch this

Thank you for the catch, and I'll update TransformLayoutTraits::UnpackedAsPython with the updated arguments. Is there somewhere that should have an additional test for round-tripping the python conversion? It looks like this UnpackedAsPython is used in Trace.as_python, but I don't see any tests that call it.

Found it. The verify_trace_roundtrip helper function verifies that the trace can be round-tripped through JSON, and that the trace can be converted to python, but doesn't verify that the trace can be round-tripped through python.

It should be corrected now, including updates to verify_trace_roundtrip to validate the python roundtrip by default. If this finds other breakages unrelated to this PR, I'll limit the python roundtrip checks to opt-in tests, with a tracking issue to resolve any other breakages.

This wasn't checked before, but was the only other issue caught by the updates to verify_trace_roundtrip.

Lunderberg · 2022-05-18T16:11:55Z

All tests passing, including tests for the UnpackAsPython. Marking the PR as ready for review.

csullivan

Very nice! I really like the new interface and refactoring under IndexMap for reuse in TE and TIR.

vinx13

Overall LGTM

Disabled while waiting for apache#11289, which was required for the `Tuple` argument.

Lunderberg · 2022-05-20T19:51:54Z

Current CI failures due to Scipy URL change, should restart CI after #11399 lands.

Lunderberg · 2022-05-24T14:36:09Z

CI failure in Docs: GPU looks like it was caused by a network timeout. Re-triggering CI.

Lunderberg · 2022-05-25T15:19:55Z

CI failure this time was due to CUDA version in Docs: GPU. Looks like the docker image was recently updating, merging main into the dev branch and retriggering CI.

Lunderberg · 2022-05-25T19:11:42Z

Retriggering CI following #11456.

Similar to apache#11269, which added this functionality to `Schedule.transform_layout`.

…11624) * [Schedule] Allowed string argument as block arg This has previously been implemented for `Schedule.transform_layout` in #11296, extending to allow for block arguments in all `Schedule` methods. This change was only made for arguments that must be a `BlockRV`. For arguments that may be either a `BlockRV` or another type (e.g. `Schedule.get_child_blocks` accepts either `BlockRV` or `LoopRV`), this sugar is not implemented, to avoid ambiguity. * [Schedule] Allowed string argument to Schedule.reindex Similar to #11269, which added this functionality to `Schedule.transform_layout`. * CI test update

…pache#11624) * [Schedule] Allowed string argument as block arg This has previously been implemented for `Schedule.transform_layout` in apache#11296, extending to allow for block arguments in all `Schedule` methods. This change was only made for arguments that must be a `BlockRV`. For arguments that may be either a `BlockRV` or another type (e.g. `Schedule.get_child_blocks` accepts either `BlockRV` or `LoopRV`), this sugar is not implemented, to avoid ambiguity. * [Schedule] Allowed string argument to Schedule.reindex Similar to apache#11269, which added this functionality to `Schedule.transform_layout`. * CI test update

Follow-up from apache#11269, which allowed schedule arguments of the buffer to be transformed to be specified as a string, or as a `tir::Buffer`. The string handling worked correctly, but the `tir::Buffer` object was handled incorrectly. This commit corrects the handling of `tir::Buffer` arguments when scheduling, and adds a unit test to validate this behavior.

…2816) Follow-up from #11269, which allowed schedule arguments of the buffer to be transformed to be specified as a string, or as a `tir::Buffer`. The string handling worked correctly, but the `tir::Buffer` object was handled incorrectly. This commit corrects the handling of `tir::Buffer` arguments when scheduling, and adds a unit test to validate this behavior.

…ache#12816) Follow-up from apache#11269, which allowed schedule arguments of the buffer to be transformed to be specified as a string, or as a `tir::Buffer`. The string handling worked correctly, but the `tir::Buffer` object was handled incorrectly. This commit corrects the handling of `tir::Buffer` arguments when scheduling, and adds a unit test to validate this behavior.

Lunderberg added 9 commits May 16, 2022 08:47

[TIR][Schedule] Added Schedule.transform_layout_sugared
e401037

[TE][TIR] Reduced duplication in TE/TIR layout transformations
7ef2312
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
bb215a1

Fix lint error
3740d75

Allow Schedule.transform_layout_sugared to set axis separators
b4195cd

Merged transform_layout_sugared functionality into transform_layout
c048d5c

Fix lint errors
7a78102

Fix lint error
3229175

Fixed docstring errors
ec66ff1

Lunderberg force-pushed the transform_layout_quality_of_life branch from 5ce8acb to ec66ff1Compare May 16, 2022 13:48

vinx13 reviewed May 17, 2022
View reviewed changes

Lunderberg added 7 commits May 17, 2022 09:31

Updated/tested TransformatLayoutTraits::UnpackedAsPython
cefee79

Merge branch 'main' into transform_layout_quality_of_life
db0f5ca

Disabled exec-used check for running trace.as_python()
5caf6f4

Updated SetAxisSeparatorTraits::UnpackedAsPython
2f63e96

Merge branch 'main' into transform_layout_quality_of_life
bd57ea2

Updated unit test that was added in merge commit
015ce32

Fixed the argument name for TensorizeTraits
accf8ff
This wasn't checked before, but was the only other issue caught by the updates to verify_trace_roundtrip.

Lunderberg marked this pull request as ready for review May 18, 2022 16:11

Lunderberg changed the title ~~[Draft][TIR][Schedule] Transform layout quality of life~~[TIR][Schedule] Transform layout quality of lifeMay 18, 2022

Lunderberg requested a review from vinx13 May 19, 2022 13:12

csullivan approved these changes May 20, 2022
View reviewed changes

vinx13 approved these changes May 20, 2022
View reviewed changes

Comment threadpython/tvm/tir/schedule/schedule.py Outdated
Comment threadpython/tvm/tir/schedule/schedule.py Outdated

Lunderberg added 2 commits May 20, 2022 12:27

Merge branch 'main' into transform_layout_quality_of_life
062a0c2

Re-enable type checks of transform_layout/set_axis_separator
99fb775
Disabled while waiting for apache#11289, which was required for the `Tuple` argument.

Lunderberg added 2 commits May 20, 2022 15:14

Merge branch 'main' into transform_layout_quality_of_life
90e0798

Updated a few additional transform_layout usages from main
a513f18

Lunderberg force-pushed the transform_layout_quality_of_life branch from a7251f2 to a513f18Compare May 23, 2022 16:30

Merge branch 'main' into transform_layout_quality_of_life
5cc3e90

Merge branch 'main' into transform_layout_quality_of_life
00e7eb6

Lunderberg force-pushed the transform_layout_quality_of_life branch from 5af3cfd to 00e7eb6Compare May 25, 2022 19:11

Merge branch 'main' into transform_layout_quality_of_life
d4b03fc

This was referenced May 25, 2022
[Hexagon] Implement avg_pool2d slice op #11417
Merged
[Hexagon] Slice op relu #11449
Merged

vinx13 merged commit 814f550 into apache:mainMay 26, 2022

Lunderberg deleted the transform_layout_quality_of_life branch May 26, 2022 11:56

Lunderberg added a commit to Lunderberg/tvm that referenced this pull request Jun 8, 2022
[Schedule] Allowed string argument to Schedule.reindex
07594c2
Similar to apache#11269, which added this functionality to `Schedule.transform_layout`.

Lunderberg mentioned this pull request Jun 8, 2022
[TIR][Schedule] Allow named block and buffer arguments in Schedule #11624
Merged

driazati mentioned this pull request Jul 14, 2022
TVM v0.9.0.rc0 Release Candidate Notes #12102
Closed

Lunderberg mentioned this pull request Sep 16, 2022
[TIR][Bugfix] Correct handling of buffer argument when scheduling #12816
Merged

Uh oh!

Conversation

Lunderberg commented May 10, 2022• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Lunderberg commented May 10, 2022

Uh oh!

vinx13 commented May 10, 2022• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Lunderberg commented May 11, 2022

Uh oh!

Lunderberg commented May 11, 2022

Uh oh!

Lunderberg commented May 12, 2022

Uh oh!

Lunderberg commented May 16, 2022

Uh oh!

vinx13May 17, 2022• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

LunderbergMay 17, 2022

Choose a reason for hiding this comment

Uh oh!

LunderbergMay 17, 2022

Choose a reason for hiding this comment

Uh oh!

LunderbergMay 17, 2022

Choose a reason for hiding this comment

Uh oh!

Lunderberg commented May 18, 2022

Uh oh!

csullivan left a comment

Choose a reason for hiding this comment

Uh oh!

vinx13 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Lunderberg commented May 20, 2022

Uh oh!

Lunderberg commented May 24, 2022

Uh oh!

Lunderberg commented May 25, 2022

Uh oh!

Lunderberg commented May 25, 2022

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Lunderberg commented May 10, 2022•
edited
Loading

vinx13 commented May 10, 2022•
edited
Loading

vinx13May 17, 2022•
edited
Loading