[TVMScript] Use TVMScript for all TIR Printing

CMakeLists.txt

@@ -288,7 +288,6 @@ tvm_file_glob(GLOB_RECURSE COMPILER_SRCS
     src/topi/*.cc
     src/driver/*.cc
     src/parser/*.cc
-    src/printer/*.cc
     src/support/*.cc
     src/script/*.cc
     )
@@ -317,6 +316,7 @@ tvm_file_glob(GLOB RELAY_BACKEND_SRCS
     )
 tvm_file_glob(GLOB_RECURSE RELAY_IR_SRCS
     src/relay/ir/*.cc
+    src/relay/printer/*.cc
     )
 tvm_file_glob(GLOB_RECURSE RELAY_QNN_SRCS
     src/relay/qnn/*.cc

include/tvm/ir/module.h

@@ -446,34 +446,6 @@ class IRModule : public ObjectRef {
   TVM_DEFINE_OBJECT_REF_COW_METHOD(IRModuleNode);
 };
 
-/*!
- * \brief Pretty print a node for debug purposes.
- *
- * \param node The node to be printed.
- * \return The text reperesentation.
- * \note This function does not show version or meta-data.
- *       Use AsText if you want to store the text.
- * \sa AsText.
- */
-TVM_DLL String PrettyPrint(const ObjectRef& node);
-
-/*!
- * \brief Render the node as a string in the text format.
- *
- * \param node The node to be rendered.
- * \param show_meta_data Whether to print meta data section.
- * \param annotate An optional callback function for attaching
- *        additional comment block to an expr.
- *
- * \note We support a limited set of IR nodes that are part of
- *       relay IR and
- *
- * \sa PrettyPrint.
- * \return The text representation.
- */
-TVM_DLL String AsText(const ObjectRef& node, bool show_meta_data = true,
-                      runtime::TypedPackedFunc<String(ObjectRef)> annotate = nullptr);
-
 namespace attr {
 
 // Following are attributes for IRModule only.

include/tvm/ir/transform.h

@@ -57,7 +57,6 @@
 #define TVM_IR_TRANSFORM_H_
 
 #include <tvm/ir/diagnostic.h>
-#include <tvm/ir/error.h>
 #include <tvm/ir/instrument.h>
 #include <tvm/ir/module.h>
 #include <tvm/runtime/container/array.h>

include/tvm/relay/base.h

@@ -120,6 +120,34 @@ class Id : public ObjectRef {
   TVM_DEFINE_OBJECT_REF_METHODS(Id, ObjectRef, IdNode);
 };
 
+/*!
+ * \brief Pretty print a node for debug purposes.
+ *
+ * \param node The node to be printed.
+ * \return The text reperesentation.
+ * \note This function does not show version or meta-data.
+ *       Use AsText if you want to store the text.
+ * \sa AsText.
+ */
+TVM_DLL String PrettyPrint(const ObjectRef& node);
+
+/*!
+ * \brief Render the node as a string in the text format.
+ *
+ * \param node The node to be rendered.
+ * \param show_meta_data Whether to print meta data section.
+ * \param annotate An optional callback function for attaching
+ *        additional comment block to an expr.
+ *
+ * \note We support a limited set of IR nodes that are part of
+ *       relay IR and
+ *
+ * \sa PrettyPrint.
+ * \return The text representation.
+ */
+TVM_DLL String AsText(const ObjectRef& node, bool show_meta_data = true,
+                      runtime::TypedPackedFunc<String(ObjectRef)> annotate = nullptr);
+
 }  // namespace relay
 }  // namespace tvm
 

include/tvm/ir/error.h → include/tvm/relay/error.h

@@ -16,13 +16,8 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-
-/*!
- * \file tvm/ir/error.h
- * \brief Utilities for error tracking and reporting.
- */
-#ifndef TVM_IR_ERROR_H_
-#define TVM_IR_ERROR_H_
+#ifndef TVM_RELAY_ERROR_H_
+#define TVM_RELAY_ERROR_H_
 
 #include <tvm/ir/module.h>
 #include <tvm/ir/span.h>
@@ -33,6 +28,7 @@
 #include <vector>
 
 namespace tvm {
+namespace relay {
 /*!
  * \brief A wrapper around std::stringstream to build error.
  *
@@ -181,5 +177,6 @@ class ErrorReporter {
   std::unordered_map<ObjectRef, GlobalVar, ObjectPtrHash, ObjectPtrEqual> node_to_gv_;
 };
 
+}  // namespace relay
 }  // namespace tvm
-#endif  // TVM_IR_ERROR_H_
+#endif  // TVM_RELAY_ERROR_H_

include/tvm/relay/expr.h

@@ -57,7 +57,6 @@ using BaseFunc = tvm::BaseFunc;
 using BaseFuncNode = tvm::BaseFuncNode;
 using GlobalVar = tvm::GlobalVar;
 using GlobalVarNode = tvm::GlobalVarNode;
-using tvm::PrettyPrint;
 
 /*!
  * \brief Constant tensor, backed by an NDArray on the cpu(0) device.

include/tvm/relay/expr_functor.h

@@ -25,9 +25,9 @@
 #ifndef TVM_RELAY_EXPR_FUNCTOR_H_
 #define TVM_RELAY_EXPR_FUNCTOR_H_
 
-#include <tvm/ir/error.h>
 #include <tvm/node/functor.h>
 #include <tvm/relay/adt.h>
+#include <tvm/relay/error.h>
 #include <tvm/relay/expr.h>
 #include <tvm/relay/function.h>
 #include <tvm/relay/op.h>

include/tvm/relay/pattern_functor.h

@@ -25,8 +25,8 @@
 #ifndef TVM_RELAY_PATTERN_FUNCTOR_H_
 #define TVM_RELAY_PATTERN_FUNCTOR_H_
 
-#include <tvm/ir/error.h>
 #include <tvm/node/functor.h>
+#include <tvm/relay/error.h>
 
 #include <string>
 #include <unordered_map>

python/tvm/ir/__init__.py

@@ -27,7 +27,6 @@
     Span,
     assert_structural_equal,
     load_json,
-    pretty_print,
     save_json,
     structural_equal,
     structural_hash,

python/tvm/ir/affine_type.py

@@ -32,7 +32,7 @@ def __ne__(self, other):
         return not self.__eq__(other)
 
     def __str__(self):
-        from tvm.ir import pretty_print  # pylint: disable=import-outside-toplevel
+        from tvm.relay import pretty_print  # pylint: disable=import-outside-toplevel
 
         return pretty_print(self)
 

python/tvm/ir/base.py

@@ -23,40 +23,9 @@
 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):
     """Base class of all IR Nodes, implements astext function."""
 
-    def astext(self, show_meta_data=True, annotate=None):
-        """Get the text format of the expression.
-
-        Parameters
-        ----------
-        show_meta_data : bool
-            Whether to include meta data section in the text
-            if there is meta data.
-
-        annotate: Optional[Object->str]
-            Optionally annotate function to provide additional
-            information in the comment block.
-
-        Returns
-        -------
-        text : str
-            The text format of the expression.
-
-        Notes
-        -----
-        The meta data section is necessary to fully parse the text format.
-        However, it can contain dumps that are big (e.g constant weights),
-        so it can be helpful to skip printing the meta data section.
-        """
-        return _ffi_api.AsText(self, show_meta_data, annotate)
-
 
 @tvm._ffi.register_object("SourceName")
 class SourceName(Object):

python/tvm/ir/expr.py

@@ -17,9 +17,9 @@
 """Common expressions data structures in the IR."""
 import tvm._ffi
 
-from .base import Node
-from . import _ffi_api
 from ..runtime import const, convert
+from . import _ffi_api
+from .base import Node
 
 
 class BaseExpr(Node):
@@ -91,6 +91,34 @@ def __call__(self, *args):
             "Do not know how to handle GlobalVar.__call__ for types {}".format(arg_types)
         )
 
+    def astext(self, show_meta_data=True, annotate=None):
+        """Get the text format of the expression.
+
+        Parameters
+        ----------
+        show_meta_data : bool
+            Whether to include meta data section in the text
+            if there is meta data.
+
+        annotate: Optional[Object->str]
+            Optionally annotate function to provide additional
+            information in the comment block.
+
+        Returns
+        -------
+        text : str
+            The text format of the expression.
+
+        Notes
+        -----
+        The meta data section is necessary to fully parse the text format.
+        However, it can contain dumps that are big (e.g constant weights),
+        so it can be helpful to skip printing the meta data section.
+        """
+        from tvm.relay import astext  # pylint: disable=import-outside-toplevel
+
+        return astext(self, show_meta_data, annotate)
+
 
 @tvm._ffi.register_object
 class Range(Node):

python/tvm/ir/module.py

@@ -287,6 +287,34 @@ def with_attr(self, attr_key, attr_value):
 
         return _ffi_api.Module_WithAttr(self, attr_key, attr_value)
 
+    def astext(self, show_meta_data=True, annotate=None):
+        """Get the text format of the expression.
+
+        Parameters
+        ----------
+        show_meta_data : bool
+            Whether to include meta data section in the text
+            if there is meta data.
+
+        annotate: Optional[Object->str]
+            Optionally annotate function to provide additional
+            information in the comment block.
+
+        Returns
+        -------
+        text : str
+            The text format of the expression.
+
+        Notes
+        -----
+        The meta data section is necessary to fully parse the text format.
+        However, it can contain dumps that are big (e.g constant weights),
+        so it can be helpful to skip printing the meta data section.
+        """
+        from tvm.relay import astext  # pylint: disable=import-outside-toplevel
+
+        return astext(self, show_meta_data, annotate)
+
     def script(
         self,
         *,

python/tvm/ir/op.py

@@ -17,8 +17,9 @@
 # pylint: disable=invalid-name
 """Primitive operators in the TVM IR."""
 import tvm._ffi
-from .expr import RelayExpr
+
 from . import _ffi_api
+from .expr import RelayExpr
 
 
 @tvm._ffi.register_object("Op")
@@ -28,6 +29,34 @@ class Op(RelayExpr):
     def __init__(self):
         raise RuntimeError("Cannot create op, use get instead")
 
+    def astext(self, show_meta_data=True, annotate=None):
+        """Get the text format of the expression.
+
+        Parameters
+        ----------
+        show_meta_data : bool
+            Whether to include meta data section in the text
+            if there is meta data.
+
+        annotate: Optional[Object->str]
+            Optionally annotate function to provide additional
+            information in the comment block.
+
+        Returns
+        -------
+        text : str
+            The text format of the expression.
+
+        Notes
+        -----
+        The meta data section is necessary to fully parse the text format.
+        However, it can contain dumps that are big (e.g constant weights),
+        so it can be helpful to skip printing the meta data section.
+        """
+        from tvm.relay import astext  # pylint: disable=import-outside-toplevel
+
+        return astext(self, show_meta_data, annotate)
+
     @staticmethod
     def get(op_name):
         """Get the Op for a given name

python/tvm/ir/tensor_type.py

@@ -56,6 +56,6 @@ def concrete_shape(self):
         return tuple(int(x) for x in self.shape)
 
     def __str__(self):
-        from tvm.ir import pretty_print  # pylint: disable=import-outside-toplevel
+        from tvm.relay import pretty_print  # pylint: disable=import-outside-toplevel
 
         return pretty_print(self)

python/tvm/micro/model_library_format.py

@@ -27,12 +27,13 @@
 
 import tvm
 from tvm.micro import get_standalone_crt_dir
+
 from .._ffi import get_global_func
 from ..contrib import utils
 from ..driver import build_module
-from ..relay.backend import executor_factory
-from ..relay.backend.name_transforms import to_c_variable_style, prefix_generated_name
 from ..relay import param_dict
+from ..relay.backend import executor_factory
+from ..relay.backend.name_transforms import prefix_generated_name, to_c_variable_style
 from ..tir import expr
 
 # This should be kept identical to runtime::symbol::tvm_module_main
@@ -528,7 +529,7 @@ def _eval_shape(param_name, buffer_shape):
         # TODO(mbs): The device type is not unique, better would be to use target.kind.name
         target_device_type = target.get_target_device_type()
         ir_mod = ir_module_by_target[target]
-        printer = get_global_func("tir.ModelLibraryFormatPrinter")(False, None, False)
+        printer = get_global_func("relay.ir.ModelLibraryFormatPrinter")(False, None, False)
         with open(src_dir / f"tir-{target_device_type}.txt", "w") as f:
             f.write(printer["print"](ir_mod))
 

python/tvm/relay/__init__.py

@@ -29,6 +29,7 @@
 from . import prelude
 from . import loops
 from . import scope_builder
+from .base import pretty_print, astext
 
 from . import transform
 from . import analysis