[Relay] Pass manager

include/tvm/relay/pass.h

@@ -55,7 +55,9 @@ namespace relay {
 
 namespace pass {
 
-// Define pass context.
+/*
+ * \brief The context of pass.
+ */
 class PassContext;
 
 /*!
@@ -82,20 +84,60 @@ class PassContextNode : public RelayNode {
 
 TVM_DEFINE_NODE_REF(PassContext, PassContextNode)
 
-class Pass;
+/*
+ * \brief The meta data of a pass.
+ *
+ * PassInfo can be extended conveniently in the future if more meta information
+ * is needed.
+ */
+class PassInfo;
 
 /*!
- * \brief PassNode is the base type of differnt types of optimization passes.
- * It is implemented by different pass subclasses at different granularity of
- * Relay nodes.
+ * \brief PassInfoNode contains meta data that will be used to help optimization
+ * and analysis.
  */
-class PassNode : public RelayNode {
+class PassInfoNode : public RelayNode {
  public:
   /*! \brief The name of an optimization/analysis pass. */
   std::string name;
+
   /*! \brief The minimal optimization level that this pass will be enabled. */
   int opt_level;
 
+  /*! \brief The passes that are required to perform the current pass. */
+  tvm::Array<tvm::Expr> required;
+
+  PassInfoNode() = default;
+
+  void VisitAttrs(tvm::AttrVisitor* v) final {
+    v->Visit("name", &name);
+    v->Visit("opt_level", &opt_level);
+    v->Visit("required", &required);
+  }
+
+  TVM_DLL static PassInfo make(std::string name, int opt_level,
+                               tvm::Array<tvm::Expr> required);
+
+  static constexpr const char* _type_key = "relay.PassInfo";
+  TVM_DECLARE_NODE_TYPE_INFO(PassInfoNode, RelayNode);
+};
+
+TVM_DEFINE_NODE_REF(PassInfo, PassInfoNode)
+
+class Pass;
+
+/*!
+ * \brief PassNode is the base type of differnt types of optimization passes.
+ * It is designed as a pure class and implemented by different pass subclasses
+ * at different granularity of Relay nodes.
+ */
+class PassNode : public RelayNode {
+ public:
+  /*
+   * \brief Get the pass information/meta data.
+   */
+  virtual PassInfo GetPassInfo() const = 0;
+
   /*!
    * \brief Set the context information for a pass.
    *
@@ -117,10 +159,7 @@ class PassNode : public RelayNode {
    */
   virtual Module operator()(const Module& mod) const = 0;
 
-  void VisitAttrs(tvm::AttrVisitor* v) override {
-    v->Visit("name", &name);
-    v->Visit("opt_level", &opt_level);
-  }
+  void VisitAttrs(tvm::AttrVisitor* v) override {}
 
   static constexpr const char* _type_key = "relay.Pass";
   TVM_DECLARE_BASE_NODE_INFO(PassNode, RelayNode);
@@ -141,40 +180,39 @@ class Pass : public NodeRef {
 /*
  * \brief Create a module pass.
  *
- * \param name The name of the module pass.
- * \param opt_level The optimization level of the module pass.
+ * \param pass_info The encapsulated pass meta data that will be used to set up
+ *        a pass.
  * \param pass_func The packed function that contains the optimization.
  *
  * \return The created module pass.
  */
 Pass CreateModulePass(
-    const std::string& name, int opt_level,
+    const PassInfo& pass_info,
     const runtime::TypedPackedFunc<Module(Module, PassContext)>& pass_func);
 
 /*
  * \brief Create a function pass.
  *
- * \param name The name of the function pass.
- * \param opt_level The optimization level of the function pass.
+ * \param pass_info The encapsulated pass meta data that will be used to set up
+ *        a pass.
  * \param pass_func The packed function that contains the optimization.
  *
  * \return The created function pass.
  */
 Pass CreateFunctionPass(
-    const std::string& name, int opt_level,
+    const PassInfo& pass_info,
     const runtime::TypedPackedFunc<Function(Function, PassContext)>& pass_func);
 /*
  * \brief Create a sequential pass.
  *
- * \param name The name of the sequential pass.
- * \param opt_level The optimization level of the sequential pass. It could be
- *        the highest opt_level of the list of passes.
+ * \param pass_info The encapsulated pass meta data that will be used to set up
+ *        a pass.
  * \param passes The optimization passes will be performed.
  * \param disabled The disabled passes.
  *
  * \return The created sequential pass.
  */
-Pass CreateSequentialPass(const std::string& name, int opt_level,
+Pass CreateSequentialPass(const PassInfo& pass_info,
                           const tvm::Array<Pass>& passes,
                           const tvm::Array<tvm::Expr>& disabled);
 

python/tvm/relay/__init__.py

@@ -95,6 +95,7 @@
 load_param_dict = param_dict.load_param_dict
 
 # Pass manager
+PassInfo = ir_pass.PassInfo
 PassContext = ir_pass.PassContext
 Pass = ir_pass.Pass
 ModulePass = ir_pass.ModulePass

python/tvm/relay/_ir_pass.pyi

@@ -8,40 +8,45 @@ class PassContext(NodeBase):
     def __init__(self):
         ...
 
-
-class Pass(NodeBase):
+class PassInfo(NodeBase):
     name = ...  # type: str
     opt_level = ... # type: int
+    required = ... # type: list
+
+    def __init__(self, name, opt_level, required)
+        # type: (str, int, list) -> None
+
+
+class Pass(NodeBase):
+    def __init__(self):
+        ...
 
 
 class ModulePass(Pass):
-    name = ...  # type: str
-    opt_level = ... # type: int
+    pass_info = ...  # type: list
     pass_func = ...  # type: Callable
 
-    def __init__(self, name, opt_level, pass_func):
-        # type: (str, int, Callable) -> None
+    def __init__(self, pass_info, pass_func):
+        # type: (list, Callable) -> None
         ...
 
 
 class FunctionPass(Pass):
-    name = ...  # type: str
-    opt_level = ... # type: int
+    pass_info = ...  # type: list
     pass_func = ...  # type: Callable
 
-    def __init__(self, name, opt_level, pass_func):
-        # type: (str, int, Callable) -> None
+    def __init__(self, pass_info, pass_func):
+        # type: (list, Callable) -> None
         ...
 
 
 class SequentialPass(Pass):
-    name = ...  # type: str
-    opt_level = ... # type: int
+    pass_info = ...  # type: list
     passes = ...  # type: list
     disabled = ... # type: list
 
-    def __init__(self, name, opt_level, passes, disabled):
-        # type: (str, int, list, list) -> None
+    def __init__(self, pass_info, passes, disabled):
+        # type: (str, list, list, list) -> None
         ...
 
 

python/tvm/relay/ir_pass.py

@@ -18,6 +18,23 @@
 from .module import Module
 
 
+@register_relay_node
+class PassInfo(RelayNode):
+    """The class that contains the meta data required by a pass. It is the
+    container of information needed by running an optimization or analysis.
+    This class can be extended by adding new members when more meta data is
+    needed.
+    """
+
+    def __init__(self, name, opt_level, required=None):
+        required = required if required else []
+        if not isinstance(required, (list, tuple)):
+            raise TypeError("required must be the list or tuple type that " +
+                            "contains a host of dependent pass namees.")
+        self.__init_handle_by_constructor__(_ir_pass.PassInfo, name, opt_level,
+                                            required)
+
+
 @register_relay_node
 class PassContext(RelayNode):
     """The basis where a Relay optimization/analysis runs on.
@@ -86,19 +103,16 @@ class ModulePass(Pass):
 
     Parameters
     ----------
-    name : str
-        The pass name.
-
-    opt_level : int
-        The optimization level of this pass.
+    pass_info : PassInfo
+        The meta data required by a module class.
 
     pass_func : Callable[(tvm.relay.Module, PassContext) -> tvm.relay.Module]
         The callback function that sketches a certain optimization.
     """
 
-    def __init__(self, name, opt_level, pass_func):
-        self.__init_handle_by_constructor__(_ir_pass.CreateModulePass, name,
-                                            opt_level, pass_func)
+    def __init__(self, pass_info, pass_func):
+        self.__init_handle_by_constructor__(_ir_pass.CreateModulePass,
+                                            pass_info, pass_func)
 
     def __call__(self, mod):
         """Execute a module pass.
@@ -122,20 +136,17 @@ class FunctionPass(Pass):
 
     Parameters
     ----------
-    name : str
-        The pass name.
-
-    opt_level : int
-        The optimization level of this pass.
+    pass_info : PassInfo
+        The meta data required by a function class.
 
     pass_func : Callable[(tvm.relay.Function, PassContext) ->
                 tvm.relay.Function]
         The callback function that sketches a certain optimization.
     """
 
-    def __init__(self, name, opt_level, pass_func):
-        self.__init_handle_by_constructor__(_ir_pass.CreateFunctionPass, name,
-                                            opt_level, pass_func)
+    def __init__(self, pass_info, pass_func):
+        self.__init_handle_by_constructor__(_ir_pass.CreateFunctionPass,
+                                            pass_info, pass_func)
 
     def __call__(self, mod):
         """Execute a function pass.
@@ -159,11 +170,8 @@ class SequentialPass(Pass):
 
     Parameters
     ----------
-    name : str
-        The pass name.
-
-    opt_level : int
-        The optimization level of this pass.
+    pass_info : PassInfo
+        The meta data required by a sequential class.
 
     passes : List[Pass]
         The pass candidates to be executed.
@@ -172,12 +180,12 @@ class SequentialPass(Pass):
         The list of passes that are disabled.
     """
 
-    def __init__(self, name, opt_level, passes, disabled=None):
+    def __init__(self, pass_info, passes, disabled=None):
         disabled = disabled if disabled else []
         if not isinstance(disabled, (list, tuple)):
             raise TypeError("disabled must be a list or tuple of pass names")
         self.__init_handle_by_constructor__(_ir_pass.CreateSequentialPass,
-                                            name, opt_level, passes, disabled)
+                                            pass_info, passes, disabled)
 
     def __call__(self, mod):
         """Execute a sequence of passes.
@@ -195,16 +203,13 @@ def __call__(self, mod):
         return _ir_pass.RunSequentialPass(self, mod)
 
 
-def create_module_pass(pass_name, opt_level, pass_func):
+def create_module_pass(pass_info, pass_func):
     """Create a module pass using a defined optimization function from Python.
 
     Parameters
     ----------
-    pass_name : str
-        The name of the pass.
-
-    opt_level : int
-        The optimization level of this pass.
+    pass_info : PassInfo
+        The meta data required by a module class.
 
     pass_func : Optional[Callable[(Module/Function, PassContext) ->
                 Module/Function]]
@@ -218,20 +223,17 @@ def create_module_pass(pass_name, opt_level, pass_func):
     if not isinstance(pass_func, (types.FunctionType, types.LambdaType)):
         raise TypeError("pass_func must be a callable for Module pass")
 
-    return _ir_pass.CreateModulePass(pass_name, opt_level, pass_func)
+    return _ir_pass.CreateModulePass(pass_info, pass_func)
 
 
-def create_function_pass(pass_name, opt_level, pass_func):
+def create_function_pass(pass_info, pass_func):
     """Create a function pass using a defined optimization function from
     Python.
 
     Parameters
     ----------
-    pass_name : str
-        The name of the pass.
-
-    opt_level : int
-        The optimization level of this pass.
+    pass_info : PassInfo
+        The meta data required by a function class.
 
     pass_func : Optional[Callable[(Module/Function, PassContext) ->
                 Module/Function]]
@@ -245,21 +247,17 @@ def create_function_pass(pass_name, opt_level, pass_func):
     if not isinstance(pass_func, (types.FunctionType, types.LambdaType)):
         raise TypeError("pass_func must be a callable for Module pass")
 
-    return _ir_pass.CreateFunctionPass(pass_name, opt_level, pass_func)
+    return _ir_pass.CreateFunctionPass(pass_info, pass_func)
 
 
-def create_sequential_pass(pass_name, opt_level, sequential_passes,
-                           disabled=None):
+def create_sequential_pass(pass_info, sequential_passes, disabled=None):
     """Create a sequential pass using a defined optimization function from
     Python.
 
     Parameters
     ----------
-    pass_name : str
-        The name of the pass.
-
-    opt_level : int
-        The optimization level of this pass.
+    pass_info : PassInfo
+        The meta data required by a sequential class.
 
     sequential_passes : Optional[List[Pass]]
         A sequence of passes candidate for optimization.
@@ -279,8 +277,8 @@ def create_sequential_pass(pass_name, opt_level, sequential_passes,
     if not isinstance(disabled, (list, tuple)):
         raise TypeError("disabled must be a list or tuple of pass names")
 
-    return _ir_pass.CreateSequentialPass(pass_name, opt_level,
-                                         sequential_passes, disabled)
+    return _ir_pass.CreateSequentialPass(pass_info, sequential_passes,
+                                         disabled)
 
 
 def post_order_visit(expr, fvisit):