[TVMScript] Introduce PrinterConfig

include/tvm/ir/expr.h

@@ -100,16 +100,7 @@ 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;
+  TVM_OBJECT_ENABLE_SCRIPT_PRINTER();
 
   static constexpr const char* _type_key = "PrimExpr";
   static constexpr const uint32_t _type_child_slots = 38;

include/tvm/ir/module.h

@@ -328,16 +328,7 @@ class IRModuleNode : public Object {
    */
   TVM_DLL std::unordered_set<String> Imports() 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
-   */
-  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;
+  TVM_OBJECT_ENABLE_SCRIPT_PRINTER();
 
   static constexpr const char* _type_key = "IRModule";
   static constexpr const bool _type_has_method_sequal_reduce = true;

include/tvm/node/repr_printer.h

@@ -24,6 +24,7 @@
 #define TVM_NODE_REPR_PRINTER_H_
 
 #include <tvm/node/functor.h>
+#include <tvm/node/script_printer.h>
 
 #include <iostream>
 #include <string>

include/tvm/node/script_printer.h

@@ -0,0 +1,105 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+/*!
+ * \file tvm/node/repr_printer.h
+ * \brief Printer class to print repr string of each AST/IR nodes.
+ */
+#ifndef TVM_NODE_SCRIPT_PRINTER_H_
+#define TVM_NODE_SCRIPT_PRINTER_H_
+
+#include <tvm/node/functor.h>
+#include <tvm/node/object_path.h>
+#include <tvm/node/reflection.h>
+#include <tvm/runtime/data_type.h>
+
+#include <iostream>
+#include <string>
+
+namespace tvm {
+
+class PrinterConfigNode : public Object {
+ public:
+  /*! \brief The prefix of IR nodes */
+  std::string ir_prefix = "I";
+  /*! \brief The prefix of TIR nodes */
+  std::string tir_prefix = "T";
+  /*! \brief The prefix of Relax nodes */
+  std::string relax_prefix = "R";
+  /*! \brief Default data type of TIR buffer */
+  DataType buffer_dtype = DataType::Float(32);
+  /*! \brief Default data type of integer literals */
+  DataType int_dtype = DataType::Int(32);
+  /*!
+   * \brief Default data type of float literals. Right now we always print out the explicit type
+   * of floating point values, so setting it to Void means we do not print without the
+   * T.float32/T.float64 wrapper.
+   */
+  DataType float_dtype = DataType::Void();
+  /*! \brief Whether or not to verbose print expressions. */
+  bool verbose_expr = false;
+  /* \brief Number of spaces used for indentation*/
+  int indent_spaces = 4;
+  /* \brief Whether to print line numbers */
+  bool print_line_numbers = false;
+  /* \brief Number of context lines to print around the underlined text */
+  int num_context_lines = -1;
+  /* \brief Object path to be underlined */
+  Optional<ObjectPath> path_to_underline = NullOpt;
+
+  void VisitAttrs(AttrVisitor* v) {
+    v->Visit("ir_prefix", &ir_prefix);
+    v->Visit("buffer_dtype", &buffer_dtype);
+    v->Visit("int_dtype", &int_dtype);
+    v->Visit("float_dtype", &float_dtype);
+    v->Visit("verbose_expr", &verbose_expr);
+    v->Visit("indent_spaces", &indent_spaces);
+    v->Visit("print_line_numbers", &print_line_numbers);
+    v->Visit("num_context_lines", &num_context_lines);
+    v->Visit("path_to_underline", &path_to_underline);
+  }
+
+  static constexpr const char* _type_key = "node.PrinterConfig";
+  TVM_DECLARE_FINAL_OBJECT_INFO(PrinterConfigNode, Object);
+};
+
+class PrinterConfig : public ObjectRef {
+ public:
+  explicit PrinterConfig(Map<String, ObjectRef> config_dict = Map<String, ObjectRef>());
+
+  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(PrinterConfig, runtime::ObjectRef,
+                                                    PrinterConfigNode);
+};
+
+/*! \brief Legacy behavior of ReprPrinter. */
+class TVMScriptPrinter {
+ public:
+  /* Convert the object to TVMScript format */
+  static std::string Script(const ObjectRef& node, const Optional<PrinterConfig>& cfg);
+  // Allow registration to be printer.
+  using FType = NodeFunctor<std::string(const ObjectRef&, const PrinterConfig&)>;
+  TVM_DLL static FType& vtable();
+};
+
+#define TVM_OBJECT_ENABLE_SCRIPT_PRINTER()                                                      \
+  std::string Script(const Optional<PrinterConfig>& config = NullOpt) const {                   \
+    return TVMScriptPrinter::Script(GetRef<ObjectRef>(this), config.value_or(PrinterConfig())); \
+  }
+
+}  // namespace tvm
+#endif  // TVM_NODE_SCRIPT_PRINTER_H_

include/tvm/script/printer/doc.h

@@ -29,8 +29,16 @@ namespace tvm {
 namespace script {
 namespace printer {
 
+// Forward declaration
 class Doc;
 
+/*!
+ * \brief Convert Doc into Python script.
+ * \param doc Doc to be converted
+ * \param cfg The configuration of the printer
+ */
+String DocToPythonScript(Doc doc, const PrinterConfig& cfg);
+
 /*!
  * \brief The base class of all Doc.
  *

include/tvm/script/printer/ir_docsifier.h

@@ -126,6 +126,8 @@ class IRDocsifierNode : public Object {
     /*! \brief The name of the variable */
     Optional<String> name;
   };
+  /*! \brief The configuration of the printer */
+  PrinterConfig cfg{nullptr};
   /*!
    * \brief The stack of frames.
    * \sa FrameNode
@@ -232,7 +234,7 @@ class IRDocsifier : public ObjectRef {
  public:
   using FType = IRDocsifierFunctor<printer::Doc, ObjectPath, IRDocsifier>;
   /*! \brief Create a IRDocsifier. */
-  IRDocsifier();
+  explicit IRDocsifier(const PrinterConfig& cfg);
   /*! \brief The registration table for IRDocsifier. */
   TVM_DLL static FType& vtable();
 

include/tvm/tir/function.h

@@ -132,16 +132,7 @@ 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;
+  TVM_OBJECT_ENABLE_SCRIPT_PRINTER();
 
   static constexpr const char* _type_key = "tir.PrimFunc";
   TVM_DECLARE_FINAL_OBJECT_INFO(PrimFuncNode, BaseFuncNode);

include/tvm/tir/stmt.h

@@ -46,16 +46,7 @@ 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;
+  TVM_OBJECT_ENABLE_SCRIPT_PRINTER();
 
   static constexpr const char* _type_key = "tir.Stmt";
   static constexpr const bool _type_has_method_sequal_reduce = true;

python/tvm/ir/expr.py

@@ -17,7 +17,7 @@
 """Common expressions data structures in the IR."""
 import tvm._ffi
 
-from ..runtime import const, convert
+from ..runtime import Scriptable, const, convert
 from . import _ffi_api
 from .base import Node
 
@@ -121,7 +121,7 @@ def astext(self, show_meta_data=True, annotate=None):
 
 
 @tvm._ffi.register_object
-class Range(Node):
+class Range(Node, Scriptable):
     """Represent a range in TVM.
 
     You do not need to create a Range explicitly.