Legalize - Use Non-recursive Rewriter. (#5296)

* Legalize - Use Non-recursive Rewriter.

* Cleanup.

src/relay/transforms/legalize.cc

@@ -35,19 +35,18 @@ namespace legalize {
 
 // Call registered FTVMLegalize of an op
 // Returns the legalized expression
-class Legalizer : public ExprMutator {
+class Legalizer : public ExprRewriter {
  public:
   explicit Legalizer(const std::string& legalize_map_attr_name)
       : legalize_map_attr_name_{legalize_map_attr_name} {}
 
-  Expr VisitExpr_(const CallNode* call_node) {
+  Expr Rewrite_(const CallNode* call_node, const Expr& post) override {
     // Get the new_call node without any changes to current call node.
-    Expr new_e = ExprMutator::VisitExpr_(call_node);
-    Call new_call = Downcast<Call>(new_e);
+    Call new_call = Downcast<Call>(post);
 
     // Check if the string is registered in the OpRegistry.
     if (!Op::HasAttr(legalize_map_attr_name_)) {
-      return new_e;
+      return post;
     }
 
     // Collect the registered legalize function.
@@ -70,19 +69,18 @@ class Legalizer : public ExprMutator {
         // Transform the op by calling the registered legalize function.
         Expr legalized_value = fop_legalize[op](call_node->attrs, call_args, types);
 
-        // Reassign new_e if the transformation succeeded.
+        // Return the new expr if the transformation succeeded.
         if (legalized_value.defined()) {
           // Check that the returned Expr from legalize is CallNode.
           const CallNode* legalized_call_node = legalized_value.as<CallNode>();
           CHECK(legalized_call_node)
               << "Can only replace the original operator with another call node";
-
-          new_e = legalized_value;
+          return legalized_value;
         }
       }
     }
 
-    return new_e;
+    return post;
   }
 
  private:
@@ -90,7 +88,8 @@ class Legalizer : public ExprMutator {
 };
 
 Expr Legalize(const Expr& expr, const std::string& legalize_map_attr_name) {
-  return Legalizer(legalize_map_attr_name).Mutate(expr);
+  auto rewriter = Legalizer(legalize_map_attr_name);
+  return PostOrderRewrite(expr, &rewriter);
 }
 
 }  // namespace legalize