[TVMScript] Migrate More to TVMScripr Printer

This PR gradually migrates more pieces of the default printing to
TVMScript printer for TIR.

include/tvm/ir/expr.h

@@ -100,6 +100,17 @@ class PrimExprNode : public BaseExprNode {
    */
   DataType dtype;
 
+  /*!
+   * \brief Returns the TVMScript format
+   * \param indent_spaces Number of spaces used for indentation
+   * \param print_line_numbers Whether to print line numbers
+   * \param num_context_lines Number of context lines to print around the underlined text
+   * \param path_to_underline Object path to be underlined
+   */
+  TVM_DLL std::string Script(int indent_spaces = 4, bool print_line_numbers = false,
+                             int num_context_lines = -1,
+                             Optional<ObjectPath> path_to_underline = NullOpt) const;
+
   static constexpr const char* _type_key = "PrimExpr";
   static constexpr const uint32_t _type_child_slots = 38;
   TVM_DECLARE_BASE_OBJECT_INFO(PrimExprNode, BaseExprNode);

include/tvm/ir/module.h

@@ -63,6 +63,26 @@ class IRModuleNode : public Object {
   parser::SourceMap source_map;
   /* \brief Additional attributes storing meta-data about the module. */
   DictAttrs attrs;
+  /*!
+   * \brief A map from string names to global variables that
+   * ensures global uniqueness.
+   */
+  Map<String, GlobalVar> global_var_map_;
+
+  /*! \brief A map from string names to global type variables (ADT names)
+   * that ensures global uniqueness.
+   */
+  Map<String, GlobalTypeVar> global_type_var_map_;
+
+  /*! \brief A map from constructor tags to constructor objects
+   * for convenient access
+   */
+  std::unordered_map<int32_t, Constructor> constructor_tag_map_;
+
+  /*! \brief The files previously imported, required to ensure
+      importing is idempotent for each module.
+   */
+  std::unordered_set<String> import_set_;
 
   /*!
    * \brief Get a module attribute.
@@ -304,15 +324,20 @@ class IRModuleNode : public Object {
   TVM_DLL void ImportFromStd(const String& path);
 
   /*!
-   * \brief Should Link Parameters into the module
-   * \return Whether the Executor is configured to execute with linked parameters (Default: false)
+   * \brief The set of imported files.
    */
-  TVM_DLL Bool ShouldLinkParameters() const;
+  TVM_DLL std::unordered_set<String> Imports() const;
 
   /*!
-   * \brief The set of imported files.
+   * \brief Returns the TVMScript format
+   * \param indent_spaces Number of spaces used for indentation
+   * \param print_line_numbers Whether to print line numbers
+   * \param num_context_lines Number of context lines to print around the underlined text
+   * \param path_to_underline Object path to be underlined
    */
-  TVM_DLL std::unordered_set<String> Imports() const;
+  TVM_DLL std::string Script(int indent_spaces = 4, bool print_line_numbers = false,
+                             int num_context_lines = -1,
+                             Optional<ObjectPath> path_to_underline = NullOpt) const;
 
   static constexpr const char* _type_key = "IRModule";
   static constexpr const bool _type_has_method_sequal_reduce = true;
@@ -322,26 +347,6 @@ class IRModuleNode : public Object {
  private:
   /*! \brief Helper function for registering a typedef's constructors */
   void RegisterConstructors(const GlobalTypeVar& var, const TypeData& type);
-
-  /*! \brief A map from string names to global variables that
-   * ensures global uniqueness.
-   */
-  Map<String, GlobalVar> global_var_map_;
-
-  /*! \brief A map from string names to global type variables (ADT names)
-   * that ensures global uniqueness.
-   */
-  Map<String, GlobalTypeVar> global_type_var_map_;
-
-  /*! \brief A map from constructor tags to constructor objects
-   * for convenient access
-   */
-  std::unordered_map<int32_t, Constructor> constructor_tag_map_;
-
-  /*! \brief The files previously imported, required to ensure
-      importing is idempotent for each module.
-   */
-  std::unordered_set<String> import_set_;
   friend class IRModule;
 };
 

include/tvm/ir/type.h

@@ -207,6 +207,25 @@ enum TypeKind : int {
   kTypeData = 6
 };
 
+/*! \brief Converts a TypeKind to a string. */
+inline String TypeKind2String(TypeKind kind) {
+  switch (kind) {
+    case TypeKind::kType:
+      return "Type";
+    case TypeKind::kShapeVar:
+      return "ShapeVar";
+    case TypeKind::kBaseType:
+      return "BaseType";
+    case TypeKind::kConstraint:
+      return "Constraint";
+    case TypeKind::kAdtHandle:
+      return "AdtHandle";
+    case TypeKind::kTypeData:
+      return "TypeData";
+  }
+  LOG(FATAL) << "ValueError: Unknown TypeKind: " << static_cast<int>(kind);
+}
+
 /*!
  * \brief Type parameter in functions.
  *

include/tvm/ir/type_functor.h

@@ -24,9 +24,9 @@
 #ifndef TVM_IR_TYPE_FUNCTOR_H_
 #define TVM_IR_TYPE_FUNCTOR_H_
 
+#include <tvm/ir/tensor_type.h>
+#include <tvm/ir/type_relation.h>
 #include <tvm/node/functor.h>
-#include <tvm/relay/adt.h>
-#include <tvm/relay/expr.h>
 
 #include <string>
 #include <utility>

include/tvm/node/repr_printer.h

@@ -26,6 +26,7 @@
 #include <tvm/node/functor.h>
 
 #include <iostream>
+#include <string>
 
 namespace tvm {
 /*! \brief A printer class to print the AST/IR nodes. */
@@ -48,6 +49,26 @@ class ReprPrinter {
   TVM_DLL static FType& vtable();
 };
 
+/*! \brief Legacy behavior of ReprPrinter. */
+class ReprLegacyPrinter {
+ public:
+  /*! \brief The output stream */
+  std::ostream& stream;
+  /*! \brief The indentation level. */
+  int indent{0};
+
+  explicit ReprLegacyPrinter(std::ostream& stream)  // NOLINT(*)
+      : stream(stream) {}
+
+  /*! \brief The node to be printed. */
+  TVM_DLL void Print(const ObjectRef& node);
+  /*! \brief Print indent to the stream */
+  TVM_DLL void PrintIndent();
+  // Allow registration to be printer.
+  using FType = NodeFunctor<void(const ObjectRef&, ReprLegacyPrinter*)>;
+  TVM_DLL static FType& vtable();
+};
+
 /*!
  * \brief Dump the node to stderr, used for debug purposes.
  * \param node The input node
@@ -70,6 +91,13 @@ inline std::ostream& operator<<(std::ostream& os, const ObjectRef& n) {  // NOLI
   ReprPrinter(os).Print(n);
   return os;
 }
+
+inline std::string AsLegacyRepr(const ObjectRef& n) {
+  std::ostringstream os;
+  ReprLegacyPrinter(os).Print(n);
+  return os.str();
+}
 }  // namespace runtime
+using runtime::AsLegacyRepr;
 }  // namespace tvm
 #endif  // TVM_NODE_REPR_PRINTER_H_

include/tvm/script/printer/ir_docsifier_functor.h

@@ -69,6 +69,9 @@ class IRDocsifierFunctor {
     if ((pf = LookupDispatchTable("", type_index)) != nullptr) {
       return (*pf)(obj, args...);
     }
+    LOG(WARNING) << "ObjectFunctor calls un-registered function on type: "
+                 << runtime::Object::TypeIndex2Key(type_index) << " (token: " << token << ")"
+                 << ". ObjectType: " << obj->GetTypeKey() << ". Object: " << obj;
     ICHECK(false) << "ObjectFunctor calls un-registered function on type: "
                   << runtime::Object::TypeIndex2Key(type_index) << " (token: " << token << ")"
                   << ". ObjectType: " << obj->GetTypeKey() << ". Object: " << obj;

include/tvm/script/printer/printer.h

@@ -55,21 +55,6 @@ struct Default {
   static bool& VerboseExpr() { return Instance()->verbose_expr; }
 };
 
-/*!
- * \brief The entry method for TVMScript printing
- * \param obj The object to be printed
- * \param indent_spaces Number of spaces used for indentation
- * \param print_line_numbers Whether to print line numbers
- * \param num_context_lines Number of context lines to print around the underlined text
- * \param path_to_underline Object path to be underlined
- * \return The TVMScript text format
- */
-String Script(ObjectRef obj,                    //
-              int indent_spaces = 4,            //
-              bool print_line_numbers = false,  //
-              int num_context_lines = -1,       //
-              Optional<ObjectPath> path_to_underline = NullOpt);
-
 /*!
  * \brief Convert Doc into Python script.
  * \param doc Doc to be converted

include/tvm/tir/expr.h

@@ -1191,9 +1191,6 @@ class Any : public PrimExpr {
   TVM_DEFINE_OBJECT_REF_COW_METHOD(AnyNode);
 };
 
-/*! \brief Legacy ReprPrint format for TIR */
-std::string LegacyTIRPrint(const ObjectRef& obj);
-
 /*
  * \brief Template function to convert Map to unordered_map
  *  Sometimes useful for API gluing when internal uses unordered_map

include/tvm/tir/function.h

@@ -132,6 +132,17 @@ class PrimFuncNode : public BaseFuncNode {
    */
   TVM_DLL FuncType func_type_annotation() const;
 
+  /*!
+   * \brief Returns the TVMScript format
+   * \param indent_spaces Number of spaces used for indentation
+   * \param print_line_numbers Whether to print line numbers
+   * \param num_context_lines Number of context lines to print around the underlined text
+   * \param path_to_underline Object path to be underlined
+   */
+  std::string Script(int indent_spaces = 4, bool print_line_numbers = false,
+                     int num_context_lines = -1,
+                     Optional<ObjectPath> path_to_underline = NullOpt) const;
+
   static constexpr const char* _type_key = "tir.PrimFunc";
   TVM_DECLARE_FINAL_OBJECT_INFO(PrimFuncNode, BaseFuncNode);
 };

include/tvm/tir/stmt.h

@@ -46,6 +46,17 @@ class StmtNode : public Object {
   StmtNode() = default;
   explicit StmtNode(Span span) : span(span) {}
 
+  /*!
+   * \brief Returns the TVMScript format
+   * \param indent_spaces Number of spaces used for indentation
+   * \param print_line_numbers Whether to print line numbers
+   * \param num_context_lines Number of context lines to print around the underlined text
+   * \param path_to_underline Object path to be underlined
+   */
+  std::string Script(int indent_spaces = 4, bool print_line_numbers = false,
+                     int num_context_lines = -1,
+                     Optional<ObjectPath> path_to_underline = NullOpt) const;
+
   static constexpr const char* _type_key = "tir.Stmt";
   static constexpr const bool _type_has_method_sequal_reduce = true;
   static constexpr const bool _type_has_method_shash_reduce = true;

python/tvm/ir/__init__.py

@@ -16,29 +16,47 @@
 # under the License.
 # pylint: disable=unused-import
 """Common data structures across all IR variants."""
-from .base import SourceName, Span, Node, EnvFunc, load_json, save_json
-from .base import structural_equal, assert_structural_equal, structural_hash
-from .type import Type, TypeKind, PrimType, PointerType, TypeVar, GlobalTypeVar, TupleType
-from .type import TypeConstraint, FuncType, IncompleteType, RelayRefType
-from .tensor_type import TensorType
-from .affine_type import TensorAffineType, TupleAffineType
-from .type_relation import TypeCall, TypeRelation
-from .expr import BaseExpr, PrimExpr, RelayExpr, GlobalVar, Range
-from .op import Op, register_op_attr, register_intrin_lowering
-from .function import CallingConv, BaseFunc
+from . import diagnostics, instrument, transform
 from .adt import Constructor, TypeData
-from .module import IRModule
+from .affine_type import TensorAffineType, TupleAffineType
 from .attrs import Attrs, DictAttrs, make_node
+from .base import (
+    EnvFunc,
+    Node,
+    SourceName,
+    Span,
+    assert_structural_equal,
+    load_json,
+    pretty_print,
+    save_json,
+    structural_equal,
+    structural_hash,
+)
 from .container import Array, Map
+from .expr import BaseExpr, GlobalVar, PrimExpr, Range, RelayExpr
+from .function import BaseFunc, CallingConv
 from .memory_pools import (
-    PoolInfo,
-    WorkspacePoolInfo,
-    ConstantPoolInfo,
-    WorkspaceMemoryPools,
     ConstantMemoryPools,
+    ConstantPoolInfo,
+    PoolInfo,
     PoolInfoProperties,
+    WorkspaceMemoryPools,
+    WorkspacePoolInfo,
 )
-
-from . import transform
-from . import instrument
-from . import diagnostics
+from .module import IRModule
+from .op import Op, register_intrin_lowering, register_op_attr
+from .tensor_type import TensorType
+from .type import (
+    FuncType,
+    GlobalTypeVar,
+    IncompleteType,
+    PointerType,
+    PrimType,
+    RelayRefType,
+    TupleType,
+    Type,
+    TypeConstraint,
+    TypeKind,
+    TypeVar,
+)
+from .type_relation import TypeCall, TypeRelation

python/tvm/ir/affine_type.py

@@ -17,8 +17,8 @@
 """Types for quantized Tensors."""
 import tvm._ffi
 
-from .base import Node
 from . import _ffi_api
+from .base import Node
 
 
 class AffineType(Node):
@@ -31,6 +31,11 @@ def __eq__(self, other):
     def __ne__(self, other):
         return not self.__eq__(other)
 
+    def __str__(self):
+        from tvm.ir import pretty_print  # pylint: disable=import-outside-toplevel
+
+        return pretty_print(self)
+
 
 @tvm._ffi.register_object("TensorAffineType")
 class TensorAffineType(AffineType):

python/tvm/ir/base.py

@@ -16,13 +16,16 @@
 # under the License.
 """Common base structures."""
 import tvm._ffi
-
 import tvm.error
 import tvm.runtime._ffi_node_api
 from tvm.runtime import Object
 
-from . import _ffi_api
-from . import json_compact
+from . import _ffi_api, json_compact
+
+
+def pretty_print(obj: Object) -> None:
+    """Pretty print the object."""
+    return _ffi_api.PrettyPrint(obj)  # type: ignore # pylint: disable=no-member
 
 
 class Node(Object):
@@ -54,9 +57,6 @@ def astext(self, show_meta_data=True, annotate=None):
         """
         return _ffi_api.AsText(self, show_meta_data, annotate)
 
-    def __str__(self):
-        return _ffi_api.PrettyPrint(self)
-
 
 @tvm._ffi.register_object("SourceName")
 class SourceName(Object):