yelite commented on code in PR #11971: URL: https://github.com/apache/tvm/pull/11971#discussion_r914103796
########## 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: These suggestions sound good to me. I made the following change: 1. Expose a `DocToPythonScript` function in the public C++/Python API (I swapped 'Python' and 'Script' to make this sound more natural IMO. Let me know if you want to use the original name you suggested in the comment) 2. Hide the base DocPrinter to a private header (as is in the previous version of code, just state here for clarity) 4. Remove `GetPythonDocPrinter` -- 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]
