[TIR][Schedule] Transform layout quality of life #11269

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

python/tvm/tir/schedule/schedule.py

-        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

python/tvm/tir/schedule/schedule.py

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

* [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

…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 ec66ff1 Compare 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 life May 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

python/tvm/tir/schedule/schedule.py Outdated Show resolved Hide resolved

python/tvm/tir/schedule/schedule.py Outdated Show resolved Hide resolved

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 a513f18 Compare 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 00e7eb6 Compare 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:main May 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

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[TIR][Schedule] Transform layout quality of life #11269

[TIR][Schedule] Transform layout quality of life #11269

Lunderberg commented May 10, 2022 •

edited

Loading

Lunderberg commented May 10, 2022

vinx13 commented May 10, 2022 •

edited

Loading

Lunderberg commented May 11, 2022

Lunderberg commented May 11, 2022

Lunderberg commented May 12, 2022

Lunderberg commented May 16, 2022

vinx13 May 17, 2022 •

edited

Loading

Lunderberg May 17, 2022

Lunderberg May 17, 2022

Lunderberg May 17, 2022

Lunderberg commented May 18, 2022

csullivan left a comment

vinx13 left a comment

Lunderberg commented May 20, 2022

Lunderberg commented May 24, 2022

Lunderberg commented May 25, 2022

Lunderberg commented May 25, 2022

[TIR][Schedule] Transform layout quality of life #11269

[TIR][Schedule] Transform layout quality of life #11269

Conversation

Lunderberg commented May 10, 2022 • edited Loading

Lunderberg commented May 10, 2022

vinx13 commented May 10, 2022 • edited Loading

Lunderberg commented May 11, 2022

Lunderberg commented May 11, 2022

Lunderberg commented May 12, 2022

Lunderberg commented May 16, 2022

vinx13 May 17, 2022 • edited Loading

Choose a reason for hiding this comment

Lunderberg May 17, 2022

Choose a reason for hiding this comment

Lunderberg May 17, 2022

Choose a reason for hiding this comment

Lunderberg May 17, 2022

Choose a reason for hiding this comment

Lunderberg commented May 18, 2022

csullivan left a comment

Choose a reason for hiding this comment

vinx13 left a comment

Choose a reason for hiding this comment

Lunderberg commented May 20, 2022

Lunderberg commented May 24, 2022

Lunderberg commented May 25, 2022

Lunderberg commented May 25, 2022

Lunderberg commented May 10, 2022 •

edited

Loading

vinx13 commented May 10, 2022 •

edited

Loading

vinx13 May 17, 2022 •

edited

Loading