[Unity][Transform] Implement RemoveUnusedParameters #16116
Conversation
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.
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_remove_unused_parameters
branch
from
59c72fb to
037d92f
Compare
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.RemoveUnusedParameters`, which is step (3) in this process.
Lunderberg
force-pushed
the
unity_transform_remove_unused_parameters
branch
from
037d92f to
3ba4640
Compare
slyubomirsky
left a comment
slyubomirsky
left a comment
There was a problem hiding this comment.
Overall, I approve of the changes. I've left a couple of questions or suggestions that we might want to address before merging.
| namespace { | ||
|
|
||
| template <typename T> | ||
| using PSet = std::unordered_set<T, ObjectPtrHash, ObjectPtrEqual>; |
There was a problem hiding this comment.
What does the P stand for?
There was a problem hiding this comment.
Here, "pointer", as the repetition of ObjectPtrHash, ObjectPtrEqual was getting a bit repetitive to read.
At some point, I want to implement std::hash for GlobalVar, tir::Var, and relax::Var. For most types, there would be ambiguity between pointer-hash and structural-hash, but the two should be equivalent in these cases.
| Function func; | ||
|
|
||
| // A mutator that updates calls at the call site. | ||
| std::function<Array<Expr>(Array<Expr>)> arg_updater; |
There was a problem hiding this comment.
The comment should probably explain how the function is meant to be used. The use of the word "mutator" can be confused for ExprMutator, so it might be better to be a little more specific.
There was a problem hiding this comment.
Good call, and updated to describe it as a function that takes as input the arguments used to call original function, and returns the arguments that should be used to call the updated function.
| for (const auto& it : callsite_updaters) { | ||
| write_ptr->Remove(it.first); | ||
| } | ||
| write_ptr->Update(new_callees); |
There was a problem hiding this comment.
Might be worth putting in a comment that this is to transfer over the functions that will be remaining unchanged. I had to think for a second to realize what this was doing.
There was a problem hiding this comment.
Good call, and I've I added a comment. This section is to remove any private functions that have unused parameters, then to replace them with updated versions of the private functions.
| CanonicalizeBindings(), | ||
| DeadCodeElimination({}), |
There was a problem hiding this comment.
Are there specific reasons to include the other passes or is it just to clean things up?
There was a problem hiding this comment.
Primarily to clean things up. The RemoveUnusedParameters doesn't leave nan floating around the way RemoveUnusedOutputs does, so these are primarily to clean things up afterward.
That said, as I'm looking at it, these really should only be applied in cases where the functions that call into the modified private functions. Therefore, moving these into the CallSiteMutator, along with a comment explaining why they are being used.
| return out | ||
|
|
||
| Expected = Before | ||
|
|
There was a problem hiding this comment.
One test case that might have interesting results would be with an inner function. I think (haven't thought it through all the way) that you might have to separately invoke the transformation on inner functions. Granted, the issue is moot if lambda lifting is used earlier.
There was a problem hiding this comment.
Hmm. Currently, it only looks for callees that are represented by a GlobalVar, so it would ignore inner functions as they are represented by a relax::Var. When an inner function is being bound, it would be recursively visited, and that would include making updates to any Call nodes encountered inside.
I'll need to add a test for it, but I think it would produce a well-formed output when a private function is called from within an inner function, but wouldn't remove unused parameters that are being used to call an inner function.
There was a problem hiding this comment.
Ah, good point, I forgot that you were looking only for global functions, so this is okay.
|
Merged |
slyubomirsky
left a comment
slyubomirsky
left a comment
There was a problem hiding this comment.
Thank you for addressing my feedback. I found the additional comments to be very helpful.
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.
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.
* [Unity] Implement RemoveUnusedParameters 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.RemoveUnusedParameters`, which is step (3) in this process. * Update based on review comments
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.
* [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
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.
2 participants
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
RemoveUnusedParameters, which removes unused parameters from the function signature, if they occur within a private function. Together withExpandTupleArgumentsimplemented in #16115, this handles partially-exposed tuple arguments for internal subroutines.