[Unity][Transform] Implement ExpandTupleArguments #16115
Merged
Conversation
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
Lunderberg
added a commit
to Lunderberg/tvm
that referenced
this pull request
Nov 13, 2023
Prior to this commit, the `FuseOps` pass explicitly tracked usage of tuple arguments, to minimize the set of arguments provided to each kernel. The additional tgracking and handling of partially-used tuples makes it difficult to follow the primary changes being made by `FuseOps`. This commit implements the same functionality in terms of the `ExpandTupleArguments` and `RemoveUnusedParameters` transforms, introduced in apache#16115 and apache#16116 respectively. By using these passes before the main `FuseOps` changes, partial tuple usage is already handled at that point. This commit is intended to minimize any changes to user-facing behavior, and so these pre-process passes are currently used internally by `FuseOps`. This may be avoided in the future by pulling this internal delegation out into a lowering pipeline.
Lunderberg
added a commit
to Lunderberg/tvm
that referenced
this pull request
Nov 13, 2023
Prior to this commit, the `FuseTIR` pass explicitly tracked usage of tuple arguments, to minimize the set of arguments provided to each kernel. The additional tgracking and handling of partially-used tuples makes it difficult to follow the primary changes being made by `FuseTIR`. This commit implements the same functionality in terms of the `ExpandTupleArguments` and `RemoveUnusedParameters` transforms, introduced in apache#16115 and apache#16116 respectively. By using these passes before the main `FuseOps` changes, partial tuple usage is already handled at that point. This commit is intended to minimize any changes to user-facing behavior, and so these pre-process passes are currently used internally by `FuseOps`. This may be avoided in the future by pulling this internal delegation out into a lowering pipeline.
Lunderberg
force-pushed
the
unity_transform_expand_tuple_arguments
branch
from
5ae7da6 to
32f2808
Compare
|
Rebased onto |
Currently, the `FuseOps` and `FuseTIR` passes have a large amount of added complexity to identify and handle partial use of tuple arguments. The handling partial use of tuples could be significantly simpler if performed in multiple steps. 1. Perform `FuseOps`. Any tuple variables that are used by the fused function are passed as-is. 2. Expand any parameters that are passed as a tuple. Any unused tensors that were included in a partially-used tuple will be converted to unused parameters. 3. Remove any unused parameters. Any unused tensors that were included in a partially-used tuple will be removed in this step. 4. Perform `FuseTIR`. No checking for tuple arguments, either partial or full, is required at this step. This PR implements `relax.transform.ExpandTupleArguments`, which is step (2) in this process.
Lunderberg
force-pushed
the
unity_transform_expand_tuple_arguments
branch
from
32f2808 to
f1d0605
Compare
slyubomirsky
approved these changes
Dec 1, 2023
Lunderberg
added a commit
to Lunderberg/tvm
that referenced
this pull request
Dec 1, 2023
Prior to this commit, the `FuseTIR` pass explicitly tracked usage of tuple arguments, to minimize the set of arguments provided to each kernel. The additional tgracking and handling of partially-used tuples makes it difficult to follow the primary changes being made by `FuseTIR`. This commit implements the same functionality in terms of the `ExpandTupleArguments` and `RemoveUnusedParameters` transforms, introduced in apache#16115 and apache#16116 respectively. By using these passes before the main `FuseOps` changes, partial tuple usage is already handled at that point. This commit is intended to minimize any changes to user-facing behavior, and so these pre-process passes are currently used internally by `FuseOps`. This may be avoided in the future by pulling this internal delegation out into a lowering pipeline.
Lunderberg
added a commit
to Lunderberg/tvm
that referenced
this pull request
Dec 1, 2023
Prior to this commit, the `FuseOps` pass explicitly tracked usage of tuple arguments, to minimize the set of arguments provided to each kernel. The additional tgracking and handling of partially-used tuples makes it difficult to follow the primary changes being made by `FuseOps`. This commit implements the same functionality in terms of the `ExpandTupleArguments` and `RemoveUnusedParameters` transforms, introduced in apache#16115 and apache#16116 respectively. By using these passes before the main `FuseOps` changes, partial tuple usage is already handled at that point. This commit is intended to minimize any changes to user-facing behavior, and so these pre-process passes are currently used internally by `FuseOps`. This may be avoided in the future by pulling this internal delegation out into a lowering pipeline.
Archermmt
pushed a commit
to Archermmt/tvm
that referenced
this pull request
Dec 18, 2023
[Unity] Implement ExpandTupleArguments transform Currently, the `FuseOps` and `FuseTIR` passes have a large amount of added complexity to identify and handle partial use of tuple arguments. The handling partial use of tuples could be significantly simpler if performed in multiple steps. 1. Perform `FuseOps`. Any tuple variables that are used by the fused function are passed as-is. 2. Expand any parameters that are passed as a tuple. Any unused tensors that were included in a partially-used tuple will be converted to unused parameters. 3. Remove any unused parameters. Any unused tensors that were included in a partially-used tuple will be removed in this step. 4. Perform `FuseTIR`. No checking for tuple arguments, either partial or full, is required at this step. This PR implements `relax.transform.ExpandTupleArguments`, which is step (2) in this process.
Lunderberg
added a commit
to Lunderberg/tvm
that referenced
this pull request
Dec 19, 2023
Prior to this commit, the `FuseTIR` pass explicitly tracked usage of tuple arguments, to minimize the set of arguments provided to each kernel. The additional tgracking and handling of partially-used tuples makes it difficult to follow the primary changes being made by `FuseTIR`. This commit implements the same functionality in terms of the `ExpandTupleArguments` and `RemoveUnusedParameters` transforms, introduced in apache#16115 and apache#16116 respectively. By using these passes before the main `FuseOps` changes, partial tuple usage is already handled at that point. This commit is intended to minimize any changes to user-facing behavior, and so these pre-process passes are currently used internally by `FuseOps`. This may be avoided in the future by pulling this internal delegation out into a lowering pipeline.
Lunderberg
added a commit
that referenced
this pull request
Jan 3, 2024
* [Unity][Transform] Extract partial-tuple-usage from FuseTIR Prior to this commit, the `FuseTIR` pass explicitly tracked usage of tuple arguments, to minimize the set of arguments provided to each kernel. The additional tgracking and handling of partially-used tuples makes it difficult to follow the primary changes being made by `FuseTIR`. This commit implements the same functionality in terms of the `ExpandTupleArguments` and `RemoveUnusedParameters` transforms, introduced in #16115 and #16116 respectively. By using these passes before the main `FuseOps` changes, partial tuple usage is already handled at that point. This commit is intended to minimize any changes to user-facing behavior, and so these pre-process passes are currently used internally by `FuseOps`. This may be avoided in the future by pulling this internal delegation out into a lowering pipeline. * Updated based on review comments * ci bump
Lunderberg
added a commit
to Lunderberg/tvm
that referenced
this pull request
Jan 3, 2024
Prior to this commit, the `FuseOps` pass explicitly tracked usage of tuple arguments, to minimize the set of arguments provided to each kernel. The additional tgracking and handling of partially-used tuples makes it difficult to follow the primary changes being made by `FuseOps`. This commit implements the same functionality in terms of the `ExpandTupleArguments` and `RemoveUnusedParameters` transforms, introduced in apache#16115 and apache#16116 respectively. By using these passes before the main `FuseOps` changes, partial tuple usage is already handled at that point. This commit is intended to minimize any changes to user-facing behavior, and so these pre-process passes are currently used internally by `FuseOps`. This may be avoided in the future by pulling this internal delegation out into a lowering pipeline.
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.
Relax functions used as subroutines may receive a tuple of arguments, but only use a subset of the arguments provided. For externally-exposed functions, this is unavoidable, as the function signature may not be updated without breaking compatibility. For internal functions, all callsites are contained within the
IRModule, and can be updated at the same time as the function signature.This PR implements
ExpandTupleArguments, which replacesR.Tupleparameters with separate parameters, when they occur in a private function.