Lunderberg commented on code in PR #15916:
URL: https://github.com/apache/tvm/pull/15916#discussion_r1372341781
##########
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:
I tend to use this construction for a couple of reasons.
* Avoid ever having a partially-initialized variable.
* Limit the scope of temporary variables that should only be used in the
initialization.
* Simplify nested if/else cases with early returns.
Effectively, the immediately-invoked lambda expression acts as a block scope
with a return type, similar to Rust's braces (e.g. The value of `{let i = 5;
i+1}` is `6`) or relax's SeqExpr.
--
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
