[Unity][Transform] Extract partial-tuple-usage from FuseTIR #16120
Conversation
Lunderberg
force-pushed
the
unity_transform_refactor_fuse_tir
branch
from
4b60b69
to
f5ac463
Compare
Lunderberg
changed the title
|
All pre-requisite PRs have landed, and this is now ready for review. |
junrushao
force-pushed
the
unity
branch
2 times, most recently
from
c95d45f
to
45eeb8c
Compare
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_refactor_fuse_tir
branch
from
d202a49
to
6656fb6
Compare
|
Rebased onto unity to avoid stale CI results. |
src/relax/transform/fuse_tir.cc
Outdated
| for (const auto& var : shape_expr->values.value()) { | ||
| ICHECK(var->IsInstance<tir::VarNode>()); | ||
| params.push_back(Downcast<tir::Var>(var)); | ||
| out->push_back(Downcast<tir::Var>(var)); | ||
| } | ||
| } else { | ||
| ICHECK(false) << "TypeError: The param type of PrimFunc is expected to be Tensor, Tuple or " |
There was a problem hiding this comment.
Nitpick: tuples are no longer accepted, so the message is inaccurate (though a tuple would have been caught above).
There was a problem hiding this comment.
Good catch, and updated to remove Tuple from the list, and add PrimValue as it is now handled.
| } | ||
|
|
||
| // Move all scalar params after buffer params. | ||
| std::stable_sort(prim_func_params.begin(), prim_func_params.end(), |
There was a problem hiding this comment.
I assume it's important for the relative ordering to be preserved (hence the stable sort), might be good to call that out.
There was a problem hiding this comment.
Good call, and updated the comment. It's mainly to make sure that the order is consistent and predictable for unit tests, not for any correctness of the output.
There was a problem hiding this comment.
The changes seem to make sense. I'm not fully versed on FuseTIR so I'm not aware of all of the broader context, but the expansion of tuples makes this code a lot a simpler and the tests all pass. I had a couple of commenting nitpicks to make, but they shouldn't be blockers. Great work.
Prior to this commit, the
FuseTIRpass explicitly tracked usage oftuple 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
ExpandTupleArgumentsandRemoveUnusedParameterstransforms,introduced in #16115 and
#16116 respectively. By using these
passes before the main
FuseOpschanges, partial tuple usage isalready 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
FuseTIR. This may be avoided in the future by pullingthis internal delegation out into a lowering pipeline.