save

apache · May 17, 2019 · f3ac8bb · f3ac8bb
1 parent 4b302cf
commit f3ac8bb
Showing 1 changed file with 33 additions and 6 deletions.
diff --git a/include/tvm/relay/pass.h b/include/tvm/relay/pass.h
@@ -507,9 +507,23 @@ inline pass::Pass FuseOpsPass(int fuse_opt_level) {
  * \return The rewritten expression.
  */
 TVM_DLL Expr ForwardRewrite(const Expr& expr,
-                    const std::string& rewrite_map_attr_name,
-                    std::function<NodeRef(const Call&)> fcontext = nullptr,
-                    std::function<Expr(const Expr&)> fmulti_ref_trigger = nullptr);
+                            const std::string& rewrite_map_attr_name,
+                            std::function<NodeRef(const Call&)> fcontext = nullptr,
+                            std::function<Expr(const Expr&)> fmulti_ref_trigger = nullptr);
+
+inline pass::Pass ForwardRewritePass(const std::string& rewrite_map_attr_name,
+                                     std::function<NodeRef(const Call&)> fcontext = nullptr,
+                                     std::function<Expr(const Expr&)>
+                                       fmulti_ref_trigger = nullptr) {
+  runtime::TypedPackedFunc<Function(Function, Module, pass::PassContext)> pass_func =
+    [=](Function f, Module m, pass::PassContext pc) {
+    return Downcast<Function>(ForwardRewrite(f,
+                                             rewrite_map_attr_name,
+                                             fcontext,
+                                             fmulti_ref_trigger));
+  };
+  return pass::CreateFunctionPass(pass_func, 1, "forward_rewrite", {});
+ }
 
 /*!
  * \brief Apply rewrite rules to rewrite the expr in post DFS order.
@@ -521,9 +535,22 @@ TVM_DLL Expr ForwardRewrite(const Expr& expr,
  * \return The rewritten expression.
  */
 TVM_DLL Expr ForwardRewrite(const Expr& expr,
-                    const FForwardRewrite& rewrite_func,
-                    std::function<NodeRef(const Call&)> fcontext = nullptr,
-                    std::function<Expr(const Expr&)> fmulti_ref_trigger = nullptr);
+                            const FForwardRewrite& rewrite_func,
+                            std::function<NodeRef(const Call&)> fcontext = nullptr,
+                            std::function<Expr(const Expr&)> fmulti_ref_trigger = nullptr);
+
+inline pass::Pass ForwardRewritePass(const FForwardRewrite& rewrite_func,
+                                     std::function<NodeRef(const Call&)> fcontext = nullptr,
+                                     std::function<Expr(const Expr&)> fmulti_ref_trigger = nullptr) {
+  runtime::TypedPackedFunc<Function(Function, Module, pass::PassContext)> pass_func =
+    [=](Function f, Module m, pass::PassContext pc) {
+    return Downcast<Function>(ForwardRewrite(f,
+                                             rewrite_func,
+                                             fcontext,
+                                             fmulti_ref_trigger));
+  };
+  return pass::CreateFunctionPass(pass_func, 1, "forward_rewrite", {});
+}
 
 /*!
  * \brief Rewrite the annotated program.