slyubomirsky commented on code in PR #15916:
URL: https://github.com/apache/tvm/pull/15916#discussion_r1372273991
##########
src/relax/transform/call_tir_rewrite.cc:
##########
@@ -111,41 +111,69 @@ class CallTIRMutator : public ExprMutator {
<< expr->struct_info_;
}
- Array<Expr> args;
- if (call->args[1].as<TupleNode>()) {
- args = Downcast<Tuple>(call->args[1])->fields;
- // for call_tir_inplace, don't reinsert in-place args, only the newly
allocated ones
- if (!is_inplace) {
- args.insert(args.end(), outs.begin(), outs.end());
- } else {
- for (size_t i = 0; i < outs.size(); i++) {
- if (inplace_attrs->inplace_indices[i].IntValue() == -1) {
- args.push_back(outs[i]);
- }
+ Expr callee = call->args[0];
+ Expr arg_tuple = call->args[1];
+ Optional<Expr> shape_tuple_of_tir_args = NullOpt;
+ if (call->args.size() > 2) {
+ shape_tuple_of_tir_args = call->args[2];
+ }
+
+ while (true) {
+ auto as_var = arg_tuple.as<Var>();
+ if (!as_var) break;
+
+ auto bound_expr = LookupBinding(as_var.value());
+ if (!bound_expr) break;
+
+ arg_tuple = bound_expr.value();
+ }
+
+ Array<Expr> args = [&]() {
+ if (auto ptr = arg_tuple.as<TupleNode>()) {
+ return ptr->fields;
+ } else if (auto ptr =
arg_tuple->struct_info_.as<TupleStructInfoNode>()) {
+ size_t n_args = ptr->fields.size();
+ Array<Expr> args;
+ for (size_t i = 0; i < n_args; i++) {
+ args.push_back(TupleGetItem(arg_tuple, i));
}
+ return args;
+ } else {
+ LOG(FATAL) << "Lowering of " << call
+ << " requires knowing how many arguments are passed to
the function. "
+ << "However, the tuple of arguments " << arg_tuple
+ << " is not itself a tuple, "
+ << "nor does its struct info " << GetStructInfo(arg_tuple)
+ << " define the number of arguments.";
}
+ }();
Review Comment:
Not that I particularly mind it, but is it preferable to use this
construction with a lambda as opposed to just assigning `args` in different
branches?
```c++
Array<Expr> args;
if (case1) {
args = ...
} else if (case2) {
args = ...
}
// etc.
```
Does it avoid an allocation or something?
--
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-unsubscr...@tvm.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
