junrushao1994 commented on code in PR #11971: URL: https://github.com/apache/tvm/pull/11971#discussion_r912381949
########## 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. Review Comment: nitpick: let's list the types line by line btw, `bool` is missing ```suggestion * Possible actual types: * - IntImm (integer or boolean) * - FloatImm * - String * - null ``` ########## src/script/printer/doc.cc: ########## @@ -0,0 +1,42 @@ +/* + * 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. + */ +#include <tvm/runtime/registry.h> +#include <tvm/script/printer/doc.h> + +namespace tvm { +namespace script { +namespace printer { + +TVM_REGISTER_NODE_TYPE(DocNode); +TVM_REGISTER_NODE_TYPE(ExprDocNode); + +LiteralDoc::LiteralDoc(ObjectRef value) { + ObjectPtr<LiteralDocNode> n = make_object<LiteralDocNode>(); + n->value = value; + this->data_ = std::move(n); +} +TVM_REGISTER_NODE_TYPE(LiteralDocNode); Review Comment: nit: let's just move all the registration logic to the same place: ```suggestion TVM_REGISTER_NODE_TYPE(ExprDocNode); LiteralDoc::LiteralDoc(ObjectRef value) { ObjectPtr<LiteralDocNode> n = make_object<LiteralDocNode>(); n->value = value; this->data_ = std::move(n); } TVM_REGISTER_NODE_TYPE(DocNode); TVM_REGISTER_NODE_TYPE(LiteralDocNode); ``` ########## 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" Review Comment: nit: ```suggestion #include <tvm/runtime/data_type.h> ``` ########## src/script/printer/python_doc_printer.cc: ########## @@ -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. + */ + +#include <tvm/runtime/registry.h> + +#include "./doc_printer.h" + +namespace tvm { +namespace script { +namespace printer { + +namespace { Review Comment: In TVM we don't use anonymous namespace very frequently because sometimes linkage could lead to some missing frames/symbols in stacktrace ########## 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)); } Review Comment: Boolean is missing ########## src/script/printer/python_doc_printer.cc: ########## @@ -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. + */ + +#include <tvm/runtime/registry.h> + +#include "./doc_printer.h" + +namespace tvm { +namespace script { +namespace printer { + +namespace { + +/*! + * \brief Print a Python literal string + * + * \param string the string to be printed + * \param out the output stream + */ +void PrintLiteralString(const String& string, std::ostringstream& out) { + // TODO(yelite): Escape and smart quote (choose ' or " automatically) + out << "\"" << string << "\""; Review Comment: See [this line](https://github.com/apache/tvm/blob/2c78daed488ea1c9f75c62587098a0c59519d8f0/src/meta_schedule/database/database_utils.cc#L75) for escape printing ########## src/script/printer/doc_printer.h: ########## @@ -0,0 +1,147 @@ +/* + * 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 <memory> +#include <ostream> +#include <string> + +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; + + /*! + * \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_ + */ + std::ostream& 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 + */ + std::ostringstream 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); Review Comment: No need to return `std::unique_ptr` - RVO is guaranteed. ```suggestion DocPrinter GetPythonDocPrinter(const DocPrinterOptions& options); ``` ########## 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) {} Review Comment: quick question: do we have the comment field for the base doc so that it could print things like: ```c++ int x; # this is doc ``` ########## python/tvm/script/printer/doc.py: ########## @@ -0,0 +1,55 @@ +# 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. +"""Doc types for TVMScript Unified Printer""" + +import tvm._ffi +from tvm.runtime import Object + +from . import _ffi_api + + +class Doc(Object): + """Base class of all Docs""" + + +class ExprDoc(Object): + """Base class of all expression Docs""" + + +LITERAL_CONSTRUCTORS = [ + (str, _ffi_api.LiteralDocStr), + (float, _ffi_api.LiteralDocFloat), + (int, _ffi_api.LiteralDocInt), +] + + +@tvm._ffi.register_object("script.printer.LiteralDoc") +class LiteralDoc(ExprDoc): + """Doc that represents literal value""" + + def __init__(self, value): + if value is None: + self.__init_handle_by_constructor__(_ffi_api.LiteralDocNone) + return + + for (cls, constructor) in LITERAL_CONSTRUCTORS: + if isinstance(value, cls): + break + else: + raise TypeError(f"Unsupported type {type(value)} for LiteralDoc") + + self.__init_handle_by_constructor__(constructor, value) Review Comment: In general, we prefer fewer indirections (even if sometimes it doesn't look the smartest): ```suggestion @tvm._ffi.register_object("script.printer.LiteralDoc") class LiteralDoc(ExprDoc): """Doc that represents literal value""" def __init__(self, value): if value is None: constructor = _ffi_api.LiteralDocNone elif isinstance(value, str): constructor = _ffi_api.LiteralDocStr elif isinstance(value, float): constructor = _ffi_api.LiteralDocFloat elif isinstance(value, int): constructor = _ffi_api.LiteralDocInt else: raise TypeError(f"Unsupported type for LiteralDoc: {type(value)}") self.__init_handle_by_constructor__(constructor, value) ``` ########## src/script/printer/python_doc_printer.cc: ########## @@ -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. + */ + +#include <tvm/runtime/registry.h> + +#include "./doc_printer.h" + +namespace tvm { +namespace script { +namespace printer { + +namespace { + +/*! + * \brief Print a Python literal string + * + * \param string the string to be printed + * \param out the output stream + */ +void PrintLiteralString(const String& string, std::ostringstream& out) { + // TODO(yelite): Escape and smart quote (choose ' or " automatically) + out << "\"" << string << "\""; +} + +/*! + * \brief Print a tvm::ir::PrimExpr as Python literal + * + * This only supports IntImm and FloatImm with size of 64 bits + * + * \param expr the PrimExpr to be printed + * \param out the output stream + */ +void PrintLiteralPrimExpr(const PrimExpr& expr, std::ostringstream& out) { + const DataType& dtype = expr->dtype; + + if (dtype == DataType::Int(64)) { + out << Downcast<IntImm>(expr)->value; + } else if (dtype == DataType::Float(64)) { + // TODO(yelite): make the float printing roundtrippable + std::ostringstream number_value; + number_value.precision(17); + number_value << Downcast<FloatImm>(expr)->value; + out << number_value.str(); + } else if (dtype == DataType::Bool()) { + out << (Downcast<IntImm>(expr)->value ? "True" : "False"); + } else { + LOG(FATAL) << "Cannot print value with dtype " << dtype << " as literal expression"; + } +} + +} // namespace + +class PythonDocPrinter : public DocPrinter { + public: + PythonDocPrinter(const DocPrinterOptions& options) : DocPrinter(options) {} + + protected: + using DocPrinter::PrintDoc; + + void PrintTypedDoc(const LiteralDoc& doc) final; +}; + +void PythonDocPrinter::PrintTypedDoc(const LiteralDoc& doc) { + const ObjectRef& value = doc->value; + if (!value.defined()) { + output_ << "None"; + } else if (const auto* expr_node = value.as<PrimExprNode>()) { + PrintLiteralPrimExpr(GetRef<PrimExpr>(expr_node), output_); + } else if (const auto* string_obj = value.as<StringObj>()) { + PrintLiteralString(GetRef<String>(string_obj), output_); + } else { + LOG(FATAL) << "Unsupported literal value type " << value->GetTypeKey(); + } +} + +std::unique_ptr<DocPrinter> GetPythonDocPrinter(const DocPrinterOptions& options) { + return std::make_unique<PythonDocPrinter>(options); +} + +TVM_REGISTER_GLOBAL("script.printer.PrintDocAsPython") + .set_body_typed([](Doc doc, int indent_spaces = 4) { + PythonDocPrinter printer({.indent_spaces = indent_spaces}); Review Comment: Designated initializers are not part of the C++ 11/14 standard IIRC, so let's avoid using it ########## src/script/printer/python_doc_printer.cc: ########## @@ -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. + */ + +#include <tvm/runtime/registry.h> + +#include "./doc_printer.h" + +namespace tvm { +namespace script { +namespace printer { + +namespace { + +/*! + * \brief Print a Python literal string + * + * \param string the string to be printed + * \param out the output stream + */ +void PrintLiteralString(const String& string, std::ostringstream& out) { + // TODO(yelite): Escape and smart quote (choose ' or " automatically) + out << "\"" << string << "\""; +} + +/*! + * \brief Print a tvm::ir::PrimExpr as Python literal + * + * This only supports IntImm and FloatImm with size of 64 bits + * + * \param expr the PrimExpr to be printed + * \param out the output stream + */ +void PrintLiteralPrimExpr(const PrimExpr& expr, std::ostringstream& out) { + const DataType& dtype = expr->dtype; + + if (dtype == DataType::Int(64)) { + out << Downcast<IntImm>(expr)->value; + } else if (dtype == DataType::Float(64)) { + // TODO(yelite): make the float printing roundtrippable + std::ostringstream number_value; + number_value.precision(17); + number_value << Downcast<FloatImm>(expr)->value; + out << number_value.str(); + } else if (dtype == DataType::Bool()) { + out << (Downcast<IntImm>(expr)->value ? "True" : "False"); + } else { + LOG(FATAL) << "Cannot print value with dtype " << dtype << " as literal expression"; + } +} Review Comment: The logic here is a bit fragmented. AFAICT, all those logic is only used in the following method: ```C++ void PrintTypedDoc(const LiteralDoc& doc) final; ``` While some of the type checking is conducted in `PrintTypedDoc`, the others are done in separate places, which makes the logic less coherent with too many lines of code. Alternatively I would prefer we merge all of those into `PrintTypeDoc` and the logic can be as succinct as: ```C++ void PythonDocPrinter::PrintTypedDoc(const LiteralDoc& doc) { const ObjectRef& value = doc->value; if (!value.defined()) { output_ << "None"; } else if (const auto* int_imm = value.as<IntImmNode>()) { if (int_imm->dtype.is_bool()) { output_ << (int_imm->value ? "True" : "False"); } else { output_ << int_imm->value; } } else if (const auto* float_imm = value.as<FloatImmNode>()) { output_.precision(17); output_ << float_imm->value; } else if (const auto* string_obj = value.as<StringObj>()) { // TODO: handle escape // See [this line](https://github.com/apache/tvm/blob/2c78daed488ea1c9f75c62587098a0c59519d8f0/src/meta_schedule/database/database_utils.cc#L75) for escape printing output_ << "\"" << string_obj->data << "\""; } else { LOG(FATAL) << "TypeError: Unsupported literal value type: " << value->GetTypeKey(); } } ``` ########## src/script/printer/python_doc_printer.cc: ########## @@ -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. + */ + +#include <tvm/runtime/registry.h> + +#include "./doc_printer.h" + +namespace tvm { +namespace script { +namespace printer { + +namespace { + +/*! + * \brief Print a Python literal string + * + * \param string the string to be printed + * \param out the output stream + */ +void PrintLiteralString(const String& string, std::ostringstream& out) { + // TODO(yelite): Escape and smart quote (choose ' or " automatically) + out << "\"" << string << "\""; +} + +/*! + * \brief Print a tvm::ir::PrimExpr as Python literal + * + * This only supports IntImm and FloatImm with size of 64 bits + * + * \param expr the PrimExpr to be printed + * \param out the output stream + */ +void PrintLiteralPrimExpr(const PrimExpr& expr, std::ostringstream& out) { + const DataType& dtype = expr->dtype; + + if (dtype == DataType::Int(64)) { + out << Downcast<IntImm>(expr)->value; + } else if (dtype == DataType::Float(64)) { + // TODO(yelite): make the float printing roundtrippable + std::ostringstream number_value; + number_value.precision(17); + number_value << Downcast<FloatImm>(expr)->value; + out << number_value.str(); + } else if (dtype == DataType::Bool()) { + out << (Downcast<IntImm>(expr)->value ? "True" : "False"); + } else { + LOG(FATAL) << "Cannot print value with dtype " << dtype << " as literal expression"; + } +} + +} // namespace + +class PythonDocPrinter : public DocPrinter { + public: + PythonDocPrinter(const DocPrinterOptions& options) : DocPrinter(options) {} + + protected: + using DocPrinter::PrintDoc; + + void PrintTypedDoc(const LiteralDoc& doc) final; +}; + +void PythonDocPrinter::PrintTypedDoc(const LiteralDoc& doc) { + const ObjectRef& value = doc->value; + if (!value.defined()) { + output_ << "None"; + } else if (const auto* expr_node = value.as<PrimExprNode>()) { + PrintLiteralPrimExpr(GetRef<PrimExpr>(expr_node), output_); + } else if (const auto* string_obj = value.as<StringObj>()) { + PrintLiteralString(GetRef<String>(string_obj), output_); + } else { + LOG(FATAL) << "Unsupported literal value type " << value->GetTypeKey(); + } +} + +std::unique_ptr<DocPrinter> GetPythonDocPrinter(const DocPrinterOptions& options) { + return std::make_unique<PythonDocPrinter>(options); +} + +TVM_REGISTER_GLOBAL("script.printer.PrintDocAsPython") Review Comment: Let's think about API consistency between python and c++: if a public API exists in c++, then it should exist also in python, and vice versa. This is a strict standard in non-legacy TVM codebase (e.g. TIR scheduling, MetaSchedule) In our particular case, I would propose: - renaming from the `PrintDocAsPython`, use `DocToScriptPython` as the public C++ / python API - `GetPythonDocPrinter` is Relatively not super useful (only used by the entry function of TVMScript printer) and there is no actual need to expose it in any header file (correct me if I was wrong) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
