slyubomirsky commented on code in PR #14394:
URL: https://github.com/apache/tvm/pull/14394#discussion_r1147086733
##########
src/relax/transform/legalize_ops.cc:
##########
@@ -85,21 +87,46 @@ class LegalizeMutator : public ExprMutator {
// Not all shape values are known
if (!std::all_of(visited_call->args.begin(), visited_call->args.end(),
- [](Expr arg) { return
KnowAllShapeValues(GetStructInfo(arg)); }) ||
+ [](Expr arg) {
+ // can happen with call_pure
+ if (arg.as<OpNode>()) {
+ return true;
+ }
+ return KnowAllShapeValues(GetStructInfo(arg));
+ }) ||
!KnowAllShapeValues(GetStructInfo(visited_call))) {
return visited_call;
}
auto op = GetRef<Op>(op_node);
+ // for call_pure, legalize the inner call
+ if (op == call_pure_op) {
+ auto inner_call = Call(call->args[0], Array<Expr>(call->args.begin() +
1, call->args.end()),
+ call->attrs, call->sinfo_args);
+ auto res = VisitExpr_(inner_call.as<CallNode>());
+ if (res.as<CallNode>()) {
+ return WrapCallPure(Downcast<Call>(res));
+ }
+ return res;
+ }
Review Comment:
One example (of many) of having to have special handling for `call_pure`.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
