[REFACTOR][IR] Polish ir/type (apache#4705)

- Use consistent constructor style to construct objects.
- Move env_func to ir as it is mainly used to construct IRs.
- Make docs consistent.

include/tvm/ir/adt.h

@@ -18,7 +18,7 @@
  */
 
 /*!
- * \file tvm/relay/adt.h
+ * \file tvm/ir/adt.h
  * \brief Algebraic data type definitions.
  *
  * We adopt relay's ADT definition as a unified class

include/tvm/node/env_func.h → include/tvm/ir/env_func.h

@@ -18,21 +18,24 @@
  */
 
 /*!
- * \file tvm/node/env_func.h
- * \brief Serializable global function.
+ * \file tvm/ir/env_func.h
+ * \brief Serializable global function used in IR.
  */
-#ifndef TVM_NODE_ENV_FUNC_H_
-#define TVM_NODE_ENV_FUNC_H_
+#ifndef TVM_IR_ENV_FUNC_H_
+#define TVM_IR_ENV_FUNC_H_
 
 #include <tvm/node/reflection.h>
 
 #include <string>
 #include <utility>
 
-
 namespace tvm {
 /*!
- * \brief Node container of EnvFunc
+ * \brief A serializable function backed by TVM's global environment.
+ *
+ * This is a wrapper to enable serializable global PackedFunc.
+ * An EnvFunc is saved by its name in the global registry
+ * under the assumption that the same function is registered during load.
  * \sa EnvFunc
  */
 class EnvFuncNode : public Object {
@@ -53,11 +56,8 @@ class EnvFuncNode : public Object {
 };
 
 /*!
- * \brief A serializable function backed by TVM's global environment.
- *
- * This is a wrapper to enable serializable global PackedFunc.
- * An EnvFunc is saved by its name in the global registry
- * under the assumption that the same function is registered during load.
+ * \brief Managed reference to EnvFuncNode.
+ * \sa EnvFuncNode
  */
 class EnvFunc : public ObjectRef {
  public:
@@ -140,4 +140,4 @@ class TypedEnvFunc<R(Args...)> : public ObjectRef {
 };
 
 }  // namespace tvm
-#endif  // TVM_NODE_ENV_FUNC_H_
+#endif  // TVM_IR_ENV_FUNC_H_

include/tvm/ir/expr.h

@@ -198,7 +198,7 @@ class GlobalVarNode : public RelayExprNode {
   /*! \brief The name of the variable, this only acts as a hint. */
   std::string name_hint;
 
-  void VisitAttrs(tvm::AttrVisitor* v) {
+  void VisitAttrs(AttrVisitor* v) {
     v->Visit("name_hint", &name_hint);
     v->Visit("span", &span);
     v->Visit("_checked_type_", &checked_type_);

include/tvm/ir/module.h

@@ -48,13 +48,13 @@ class IRModule;
 class IRModuleNode : public Object {
  public:
   /*! \brief A map from ids to all global functions. */
-  tvm::Map<GlobalVar, BaseFunc> functions;
+  Map<GlobalVar, BaseFunc> functions;
   /*! \brief A map from global type vars to ADT type data. */
-  tvm::Map<GlobalTypeVar, TypeData> type_definitions;
+  Map<GlobalTypeVar, TypeData> type_definitions;
 
   IRModuleNode() {}
 
-  void VisitAttrs(tvm::AttrVisitor* v) {
+  void VisitAttrs(AttrVisitor* v) {
     v->Visit("functions", &functions);
     v->Visit("type_definitions", &type_definitions);
     v->Visit("global_var_map_", &global_var_map_);
@@ -146,7 +146,7 @@ class IRModuleNode : public Object {
    * \brief Collect all global vars defined in this module.
    * \returns An array of global vars
    */
-  TVM_DLL tvm::Array<GlobalVar> GetGlobalVars() const;
+  TVM_DLL Array<GlobalVar> GetGlobalVars() const;
 
   /*!
    * \brief Look up a global function by its name.
@@ -159,7 +159,7 @@ class IRModuleNode : public Object {
    * \brief Collect all global type vars defined in this module.
    * \returns An array of global type vars
    */
-  TVM_DLL tvm::Array<GlobalTypeVar> GetGlobalTypeVars() const;
+  TVM_DLL Array<GlobalTypeVar> GetGlobalTypeVars() const;
 
   /*!
    * \brief Look up a global function by its variable.
@@ -235,12 +235,12 @@ class IRModuleNode : public Object {
   /*! \brief A map from string names to global variables that
    * ensures global uniqueness.
    */
-  tvm::Map<std::string, GlobalVar> global_var_map_;
+  Map<std::string, GlobalVar> global_var_map_;
 
   /*! \brief A map from string names to global type variables (ADT names)
    * that ensures global uniqueness.
    */
-  tvm::Map<std::string, GlobalTypeVar> global_type_var_map_;
+  Map<std::string, GlobalTypeVar> global_type_var_map_;
 
   /*! \brief A map from constructor tags to constructor objects
    * for convenient access
@@ -266,8 +266,8 @@ class IRModule : public ObjectRef {
    * \param type_definitions Type definitions in the module.
    * \param import_set Set of imported files in the module
    */
-  TVM_DLL explicit IRModule(tvm::Map<GlobalVar, BaseFunc> functions,
-                            tvm::Map<GlobalTypeVar, TypeData> type_definitions = {},
+  TVM_DLL explicit IRModule(Map<GlobalVar, BaseFunc> functions,
+                            Map<GlobalTypeVar, TypeData> type_definitions = {},
                             std::unordered_set<std::string> import_set = {});
   /*! \brief default constructor */
   IRModule() {}
@@ -296,8 +296,8 @@ class IRModule : public ObjectRef {
    */
   TVM_DLL static IRModule FromExpr(
     const RelayExpr& expr,
-    const tvm::Map<GlobalVar, BaseFunc>& global_funcs = {},
-    const tvm::Map<GlobalTypeVar, TypeData>& type_definitions = {});
+    const Map<GlobalVar, BaseFunc>& global_funcs = {},
+    const Map<GlobalTypeVar, TypeData>& type_definitions = {});
 
   /*!
    * \brief Parse text format source file into an IRModule.

include/tvm/ir/op.h

@@ -91,7 +91,7 @@ class OpNode : public RelayExprNode {
    */
   int32_t support_level = 10;
 
-  void VisitAttrs(tvm::AttrVisitor* v) {
+  void VisitAttrs(AttrVisitor* v) {
     v->Visit("name", &name);
     v->Visit("op_type", &op_type);
     v->Visit("description", &description);
@@ -476,15 +476,15 @@ inline OpRegistry& OpRegistry::add_type_rel(
   std::string input_name_prefix = "in";
   for (int i = 0; i < get()->num_inputs; i++) {
     auto name = input_name_prefix + std::to_string(i);
-    auto param = TypeVarNode::make(name, TypeKind::kType);
+    auto param = TypeVar(name, TypeKind::kType);
     type_params.push_back(param);
     arg_types.push_back(param);
   }
 
   Array<Type> ty_call_args = arg_types;
 
   // Add output type.
-  auto out_param = TypeVarNode::make("out", TypeKind::kType);
+  auto out_param = TypeVar("out", TypeKind::kType);
   type_params.push_back(out_param);
   // this will trigger copy on write.
   ty_call_args.push_back(out_param);
@@ -498,13 +498,13 @@ inline OpRegistry& OpRegistry::add_type_rel(
   // A common example is sum(x, axis), where the choice of axis
   // can affect the type of the function.
   TypeConstraint type_rel =
-      TypeRelationNode::make(env_type_rel_func,
+      TypeRelation(env_type_rel_func,
                              ty_call_args,
                              arg_types.size(),
                              Attrs());
 
   auto func_type =
-      FuncTypeNode::make(arg_types, out_param, type_params, {type_rel});
+      FuncType(arg_types, out_param, type_params, {type_rel});
 
   get()->op_type = func_type;
 

include/tvm/ir/transform.h

@@ -84,13 +84,13 @@ class PassContextNode : public Object {
   int fallback_device{static_cast<int>(kDLCPU)};
 
   /*! \brief The list of required passes. */
-  tvm::Array<tvm::PrimExpr> required_pass;
+  Array<PrimExpr> required_pass;
   /*! \brief The list of disabled passes. */
-  tvm::Array<tvm::PrimExpr> disabled_pass;
+  Array<PrimExpr> disabled_pass;
 
   PassContextNode() = default;
 
-  void VisitAttrs(tvm::AttrVisitor* v) {
+  void VisitAttrs(AttrVisitor* v) {
     v->Visit("opt_level", &opt_level);
     v->Visit("fallback_device", &fallback_device);
     v->Visit("required_pass", &required_pass);
@@ -118,7 +118,7 @@ class PassContextNode : public Object {
 class PassContext : public ObjectRef {
  public:
   PassContext() {}
-  explicit PassContext(ObjectPtr<::tvm::Object> n) : ObjectRef(n) {}
+  explicit PassContext(ObjectPtr<Object> n) : ObjectRef(n) {}
   /*!
    * \brief const accessor.
    * \return const access pointer.
@@ -158,7 +158,7 @@ class PassContext : public ObjectRef {
 
   // Classes to get the Python `with` like syntax.
   friend class Internal;
-  friend class tvm::With<PassContext>;
+  friend class With<PassContext>;
 };
 
 /*!
@@ -174,11 +174,11 @@ class PassInfoNode : public Object {
   std::string name;
 
   /*! \brief The passes that are required to perform the current pass. */
-  tvm::Array<tvm::PrimExpr> required;
+  Array<PrimExpr> required;
 
   PassInfoNode() = default;
 
-  void VisitAttrs(tvm::AttrVisitor* v) {
+  void VisitAttrs(AttrVisitor* v) {
     v->Visit("opt_level", &opt_level);
     v->Visit("name", &name);
     v->Visit("required", &required);
@@ -202,7 +202,7 @@ class PassInfo : public ObjectRef {
    */
   TVM_DLL PassInfo(int opt_level,
                    std::string name,
-                   tvm::Array<tvm::PrimExpr> required);
+                   Array<PrimExpr> required);
 
   TVM_DEFINE_OBJECT_REF_METHODS(PassInfo, ObjectRef, PassInfoNode);
 };
@@ -241,7 +241,7 @@ class PassNode : public Object {
   virtual IRModule operator()(const IRModule& mod,
                               const PassContext& pass_ctx) const = 0;
 
-  void VisitAttrs(tvm::AttrVisitor* v) {}
+  void VisitAttrs(AttrVisitor* v) {}
 
   static constexpr const char* _type_key = "relay.Pass";
   TVM_DECLARE_BASE_OBJECT_INFO(PassNode, Object);
@@ -289,7 +289,7 @@ class Sequential : public Pass {
    * \param passes The passes to apply.
    * \param pass_info The pass metadata.
    */
-  TVM_DLL Sequential(tvm::Array<Pass> passes, PassInfo pass_info);
+  TVM_DLL Sequential(Array<Pass> passes, PassInfo pass_info);
 
   /*!
    * \brief The constructor of `Sequential`.
@@ -299,10 +299,10 @@ class Sequential : public Pass {
    *        This allows users to only provide a list of passes and execute them
    *        under a given context.
    */
-  TVM_DLL Sequential(tvm::Array<Pass> passes, std::string name = "sequential");
+  TVM_DLL Sequential(Array<Pass> passes, std::string name = "sequential");
 
   Sequential() = default;
-  explicit Sequential(tvm::ObjectPtr<::tvm::Object> n) : Pass(n) {}
+  explicit Sequential(ObjectPtr<Object> n) : Pass(n) {}
 
   const SequentialNode* operator->() const;
   using ContainerType = Sequential;
@@ -322,7 +322,7 @@ Pass CreateModulePass(
     const runtime::TypedPackedFunc<IRModule(IRModule, PassContext)>& pass_func,
     int opt_level,
     const std::string& name,
-    const tvm::Array<tvm::PrimExpr>& required);
+    const Array<PrimExpr>& required);
 
 }  // namespace transform
 }  // namespace tvm