You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by "Lunderberg (via GitHub)" <gi...@apache.org> on 2024/02/16 17:34:01 UTC
[PR] [Unity][Transform] Handle dynamic shapes in CombineParallelMatmul [tvm]
Lunderberg opened a new pull request, #16591:
URL: https://github.com/apache/tvm/pull/16591
Prior to this commit, if the weight of a matmul a dynamic shape, and that matmul is being combined with the `CombineParallelMatmul` pass, it could cause a segfault when `dim.as<IntImmNode>()` returns a null pointer.
This commit adds explicit test cases for these dynamic shapes, and updates `CombineParallelMatmul` to handle the dynamic shapes.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Handle dynamic shapes in CombineParallelMatmul [tvm]
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on code in PR #16591:
URL: https://github.com/apache/tvm/pull/16591#discussion_r1500762292
##########
src/relax/transform/combine_parallel_matmul.cc:
##########
@@ -140,40 +149,68 @@ runtime::TypedPackedFunc<Map<Var, Expr>(Map<DFPattern, Var>, Map<Var, Expr>)> Ge
for (const auto& [rhs_dim, indices] : GroupShapes(rhs_shapes)) {
if (indices.size() == 1 || !batch_dims_compatible(rhs_dim, indices, rhs_shapes)) continue;
- auto inp = matchings[patterns.input];
+ auto lhs = matchings[patterns.input];
+
+ const auto& patterns_to_replace = [&patterns, &branch_info]() {
+ if (branch_info.activation) return patterns.activation;
+ if (branch_info.bias_dim) return patterns.bias_add;
+ return patterns.matmul;
+ }();
- Array<Var> rhs, bias;
- for (auto ind : indices) {
- rhs.push_back(matchings[patterns.rhs[ind]]);
- if (branch_info.bias_dim) {
- ICHECK(matchings.count(patterns.bias[ind]));
- bias.push_back(matchings[patterns.bias[ind]]);
+ std::vector<SplitInfo> splits;
+ for (auto index : indices) {
+ Var rhs = matchings[patterns.rhs[index]];
+ Optional<Var> bias = NullOpt;
+ if (branch_info.bias_dim.has_value()) {
+ bias = matchings[patterns.bias[index]];
}
+ PrimExpr split_size = GetTensorSInfo(rhs)->GetShape().value()[rhs_dim - 1];
+ DFPattern pattern_to_replace = patterns_to_replace[index];
+ splits.push_back(SplitInfo{rhs, bias, split_size, pattern_to_replace});
+ }
+ // At most one dynamic output shape can be part of the combined
+ // matmul, and it must be the last item in the split. Use
+ // `std::stable_sort` instead of `std::sort` to maintain a
+ // consistent order for all static shapes, and to consistently
+ // select the same dynamic weight to participate.
+ auto is_dynamic_split = [](const SplitInfo& split) -> bool {
+ return !split.split_size->IsInstance<IntImmNode>();
+ };
+ std::stable_sort(splits.begin(), splits.end(),
+ [&is_dynamic_split](const auto& a, const auto& b) {
+ return is_dynamic_split(a) < is_dynamic_split(b);
Review Comment:
Thank you. I went back and forth on whether this was reasonably clever, or too clever, and I think I like it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Handle dynamic shapes in CombineParallelMatmul [tvm]
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg merged PR #16591:
URL: https://github.com/apache/tvm/pull/16591
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Handle dynamic shapes in CombineParallelMatmul [tvm]
Posted by "slyubomirsky (via GitHub)" <gi...@apache.org>.
slyubomirsky commented on code in PR #16591:
URL: https://github.com/apache/tvm/pull/16591#discussion_r1500182313
##########
src/relax/transform/combine_parallel_matmul.cc:
##########
@@ -140,40 +149,68 @@ runtime::TypedPackedFunc<Map<Var, Expr>(Map<DFPattern, Var>, Map<Var, Expr>)> Ge
for (const auto& [rhs_dim, indices] : GroupShapes(rhs_shapes)) {
if (indices.size() == 1 || !batch_dims_compatible(rhs_dim, indices, rhs_shapes)) continue;
- auto inp = matchings[patterns.input];
+ auto lhs = matchings[patterns.input];
+
+ const auto& patterns_to_replace = [&patterns, &branch_info]() {
+ if (branch_info.activation) return patterns.activation;
+ if (branch_info.bias_dim) return patterns.bias_add;
+ return patterns.matmul;
+ }();
- Array<Var> rhs, bias;
- for (auto ind : indices) {
- rhs.push_back(matchings[patterns.rhs[ind]]);
- if (branch_info.bias_dim) {
- ICHECK(matchings.count(patterns.bias[ind]));
- bias.push_back(matchings[patterns.bias[ind]]);
+ std::vector<SplitInfo> splits;
+ for (auto index : indices) {
+ Var rhs = matchings[patterns.rhs[index]];
+ Optional<Var> bias = NullOpt;
+ if (branch_info.bias_dim.has_value()) {
+ bias = matchings[patterns.bias[index]];
}
+ PrimExpr split_size = GetTensorSInfo(rhs)->GetShape().value()[rhs_dim - 1];
+ DFPattern pattern_to_replace = patterns_to_replace[index];
+ splits.push_back(SplitInfo{rhs, bias, split_size, pattern_to_replace});
+ }
+ // At most one dynamic output shape can be part of the combined
+ // matmul, and it must be the last item in the split. Use
+ // `std::stable_sort` instead of `std::sort` to maintain a
+ // consistent order for all static shapes, and to consistently
+ // select the same dynamic weight to participate.
+ auto is_dynamic_split = [](const SplitInfo& split) -> bool {
+ return !split.split_size->IsInstance<IntImmNode>();
+ };
+ std::stable_sort(splits.begin(), splits.end(),
+ [&is_dynamic_split](const auto& a, const auto& b) {
+ return is_dynamic_split(a) < is_dynamic_split(b);
Review Comment:
Very clever to use a comparison between bools.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org