[TVMScript] Doc Base Class & DocPrinter Scaffolding

CMakeLists.txt

@@ -281,6 +281,7 @@ tvm_file_glob(GLOB_RECURSE COMPILER_SRCS
     src/parser/*.cc
     src/printer/*.cc
     src/support/*.cc
+    src/script/*.cc
     )
 
 tvm_file_glob(GLOB CODEGEN_SRCS

include/tvm/script/printer/doc.h

@@ -0,0 +1,157 @@
+/*
+ * 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.
+ */
+#ifndef TVM_SCRIPT_PRINTER_DOC_H_
+#define TVM_SCRIPT_PRINTER_DOC_H_
+
+#include <tvm/ir/expr.h>
+#include <tvm/node/node.h>
+
+#include "tvm/runtime/data_type.h"
+
+namespace tvm {
+namespace script {
+namespace printer {
+
+/*!
+ * \brief The base class of all Doc.
+ *
+ * Doc is an intermediate representation between IR from TVM
+ * and the TVMScript code.
+ * During printing, IR graph is first translated into Doc tree,
+ * then the Doc tree is translated to the target language in
+ * text format.
+ *
+ * \sa Doc
+ */
+class DocNode : public Object {
+ public:
+  void VisitAttrs(AttrVisitor* v) {}
+
+  static constexpr const char* _type_key = "script.printer.Doc";
+  TVM_DECLARE_BASE_OBJECT_INFO(DocNode, Object);
+
+ public:
+  virtual ~DocNode() = default;
+};
+
+/*!
+ * \brief Reference type of DocNode.
+ *
+ * \sa DocNode
+ */
+class Doc : public ObjectRef {
+ protected:
+  Doc() = default;
+
+ public:
+  virtual ~Doc() = default;
+  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(Doc, ObjectRef, DocNode);
+};
+
+/*!
+ * \brief The base class of expression doc.
+ *
+ * \sa ExprDoc
+ */
+class ExprDocNode : public DocNode {
+ public:
+  void VisitAttrs(AttrVisitor* v) { DocNode::VisitAttrs(v); }
+
+  static constexpr const char* _type_key = "script.printer.ExprDoc";
+  TVM_DECLARE_BASE_OBJECT_INFO(ExprDocNode, DocNode);
+};
+
+/*!
+ * \brief Reference type of ExprDocNode.
+ *
+ * \sa ExprDocNode
+ */
+class ExprDoc : public Doc {
+ protected:
+  ExprDoc() = default;
+
+ public:
+  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(ExprDoc, Doc, ExprDocNode);
+};
+
+/*!
+ * \brief Doc that represents literal value.
+ *
+ * \sa LiteralDoc
+ */
+class LiteralDocNode : public ExprDocNode {
+ public:
+  /*!
+   * \brief the internal representation of the literal value.
+   *
+   * The actual type is union of IntImm, FloatImm and String, or a
+   * null ObjectRef.
+   */
+  ObjectRef value;
+
+  void VisitAttrs(AttrVisitor* v) {
+    ExprDocNode::VisitAttrs(v);
+    v->Visit("value", &value);
+  }
+
+  static constexpr const char* _type_key = "script.printer.LiteralDoc";
+  TVM_DECLARE_FINAL_OBJECT_INFO(LiteralDocNode, ExprDocNode);
+};
+
+/*!
+ * \brief Reference type of LiteralDocNode.
+ *
+ * \sa LiteralDocNode
+ */
+class LiteralDoc : public ExprDoc {
+ protected:
+  explicit LiteralDoc(ObjectRef value);
+
+ public:
+  /*!
+   * \brief Create a LiteralDoc to represent None/null/empty value.
+   */
+  static LiteralDoc None() { return LiteralDoc(ObjectRef(nullptr)); }
+
+  /*!
+   * \brief Create a LiteralDoc to represent integer.
+   * \param v The integer value.
+   */
+  static LiteralDoc Int(int v) { return LiteralDoc(IntImm(DataType::Int(64), v)); }
+
+  /*!
+   * \brief Create a LiteralDoc to represent float.
+   * \param v The float value.
+   */
+  static LiteralDoc Float(double v) { return LiteralDoc(FloatImm(DataType::Float(64), v)); }
+
+  /*!
+   * \brief Create a LiteralDoc to represent string.
+   * \param v The string value.
+   */
+  static LiteralDoc Str(const String& v) { return LiteralDoc(v); }
+
+  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(LiteralDoc, ExprDoc, LiteralDocNode);
+};
+
+}  // namespace printer
+}  // namespace script
+}  // namespace tvm
+
+#endif  // TVM_SCRIPT_PRINTER_DOC_H_

include/tvm/script/printer/doc_printer.h

@@ -0,0 +1,148 @@
+/*
+ * 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.
+ */
+#ifndef TVM_SCRIPT_PRINTER_DOC_PRINTER_H_
+#define TVM_SCRIPT_PRINTER_DOC_PRINTER_H_
+
+#include <tvm/script/printer/doc.h>
+
+#include <string>
+#include <memory>
+
+namespace tvm {
+namespace script {
+namespace printer {
+
+/*!
+ * \brief Configurable options for DocPrinter
+ *
+ * \sa DocPrinter
+ */
+struct DocPrinterOptions {
+  int indent_spaces = 4;
+};
+
+/*!
+ * \brief DocPrinter is responsible for printing Doc tree into text format
+ * \details This is the base class for translating Doc into string.
+ *          Each target language needs to have its subclass of DocPrinter
+ *          to define the actual logic of printing Doc.
+ *
+ * \sa Doc
+ */
+class DocPrinter {
+ public:
+  /*!
+   * \brief The constructor of DocPrinter
+   *
+   * \param options the option for printer
+   */
+  explicit DocPrinter(const DocPrinterOptions& options);
+  virtual ~DocPrinter() = default;
+
+  /*!
+   * \brief Append a doc into the final content
+   *
+   * \param doc the Doc to be printed
+   *
+   * \sa GetString
+   */
+  void Append(const Doc& doc);
+
+  /*!
+   * \brief Get the printed string of all Doc appended
+   *
+   * The content of each Doc in the returned string will
+   * appear in the same order as they are appended.
+   *
+   * \sa Append
+   */
+  String GetString() const;
+
+ protected:
+  /*!
+   * \brief Get the printed string
+   *
+   * It will dispatch to the PrintTypedDoc method based on
+   * the actual type of Doc.
+   *
+   * \sa PrintTypedDoc
+   */
+  void PrintDoc(const Doc& doc);
+
+  /*!
+   * \brief Virtual method to print a LiteralDoc
+   */
+  virtual void PrintTypedDoc(const LiteralDoc& doc) = 0;
+
+  using OutputStream = std::ostringstream;
+
+  /*!
+   * \brief Increase the indent level of any content to be
+   *        printed after this call
+   */
+  void IncreaseIndent() { indent_ += options_.indent_spaces; }
+
+  /*!
+   * \brief Decrease the indent level of any content to be
+   *        printed after this call
+   */
+  void DecreaseIndent() { indent_ -= options_.indent_spaces; }
+
+  /*!
+   * \brief Add a new line into the output stream
+   *
+   * \sa output_
+   */
+  OutputStream& NewLine() {
+    output_ << "\n";
+    output_ << std::string(indent_, ' ');
+    return output_;
+  }
+
+  /*!
+   * \brief The output stream of printer
+   *
+   * All printed content will be stored in this stream and returned
+   * when GetString is called.
+   *
+   * \sa GetString
+   */
+  OutputStream output_;
+
+ private:
+  /*! \brief the printer options */
+  DocPrinterOptions options_;
+
+  /*! \brief the current level of indent */
+  int indent_ = 0;
+};
+
+/*!
+ * \brief Get a doc printer to print Doc into Python code
+ *
+ * \param options the option for printer
+ * \return A pointer to the printer
+ */
+std::unique_ptr<DocPrinter> GetPythonDocPrinter(const DocPrinterOptions& options);
+
+}  // namespace printer
+}  // namespace script
+}  // namespace tvm
+
+#endif  // TVM_SCRIPT_PRINTER_DOC_PRINTER_H_

python/tvm/script/printer/__init__.py

@@ -0,0 +1,26 @@
+# 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.
+"""
+TVMScript Unified Printer
+
+This package provides a set of APIs to print supported TVM IR into TVMScript
+in a roundtrippable way.
+
+https://github.com/apache/tvm-rfcs/blob/main/rfcs/0074-tvmscript-unified-printer.md
+"""
+
+from . import _ffi_api

python/tvm/script/printer/_ffi_api.py

@@ -0,0 +1,20 @@
+# 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.
+"""FFI APIs for tvm.script.printer"""
+import tvm._ffi
+
+tvm._ffi._init_api("script.printer", __name__)