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) { ...@@ -83,27 +83,27 @@ Expr ExprMutator::VisitExpr_(const TupleNode* op) {
Expr ExprMutator::VisitExpr_(const FunctionNode* op) { Expr ExprMutator::VisitExpr_(const FunctionNode* op) {
tvm::Array<TypeVar> ty_params; tvm::Array<TypeVar> ty_params;
bool all_ty_params_changed = true; bool all_ty_params_unchanged = true;
for (auto ty_param : op->type_params) { for (auto ty_param : op->type_params) {
TypeVar new_ty_param = Downcast<TypeVar>(VisitType(ty_param)); TypeVar new_ty_param = Downcast<TypeVar>(VisitType(ty_param));
ty_params.push_back(new_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; tvm::Array<Var> params;
bool all_params_changed = true; bool all_params_unchanged = true;
for (auto param : op->params) { for (auto param : op->params) {
Var new_param = Downcast<Var>(this->Mutate(param)); Var new_param = Downcast<Var>(this->Mutate(param));
params.push_back(new_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 ret_type = this->VisitType(op->ret_type);
auto body = this->Mutate(op->body); auto body = this->Mutate(op->body);
if (ty_params.same_as(op->type_params) && if (all_ty_params_unchanged &&
params.same_as(op->params) && all_params_unchanged &&
ret_type.same_as(op->ret_type) && ret_type.same_as(op->ret_type) &&
body.same_as(op->body)) { body.same_as(op->body)) {
return GetRef<Expr>(op); 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