Commit 4b530561 by lixiaoquan Committed by Tianqi Chen

[RELAY] Avoid unnecessarily reconstructing FunctionNode. (#3047)

parent 85a3ea08
......@@ -83,27 +83,27 @@ Expr ExprMutator::VisitExpr_(const TupleNode* op) {
Expr ExprMutator::VisitExpr_(const FunctionNode* op) {
tvm::Array<TypeVar> ty_params;
bool all_ty_params_changed = true;
bool all_ty_params_unchanged = true;
for (auto ty_param : op->type_params) {
TypeVar new_ty_param = Downcast<TypeVar>(VisitType(ty_param));
ty_params.push_back(new_ty_param);
all_ty_params_changed &= new_ty_param.same_as(ty_param);
all_ty_params_unchanged &= new_ty_param.same_as(ty_param);
}
tvm::Array<Var> params;
bool all_params_changed = true;
bool all_params_unchanged = true;
for (auto param : op->params) {
Var new_param = Downcast<Var>(this->Mutate(param));
params.push_back(new_param);
all_params_changed &= param.same_as(new_param);
all_params_unchanged &= param.same_as(new_param);
}
auto ret_type = this->VisitType(op->ret_type);
auto body = this->Mutate(op->body);
if (ty_params.same_as(op->type_params) &&
params.same_as(op->params) &&
if (all_ty_params_unchanged &&
all_params_unchanged &&
ret_type.same_as(op->ret_type) &&
body.same_as(op->body)) {
return GetRef<Expr>(op);
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment