/* * 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. */ /*! * Copyright (c) 2019 by Contributors * \file tvm/relay/doc.h * \brief Doc ADT used for pretty printing. * Based on Section 1 of * https://homepages.inf.ed.ac.uk/wadler/papers/prettier/prettier.pdf, but with * a vector instead of an implicitly linked list. */ #ifndef TVM_RELAY_IR_DOC_H_ #define TVM_RELAY_IR_DOC_H_ #include <tvm/relay/expr.h> #include <memory> #include <string> #include <vector> namespace tvm { namespace relay { // Doc Atom ADT struct DocAtomNode { virtual ~DocAtomNode() = default; }; using DocAtom = std::shared_ptr<DocAtomNode>; struct TextNode : DocAtomNode { std::string str; explicit TextNode(const std::string& str) : str(str) {} }; struct LineNode : DocAtomNode { int indent; explicit LineNode(int indent) : indent(indent) {} }; // Doc is a stream-like interface class Doc { public: Doc() {} explicit Doc(const std::string& str); // Append right to this. Doc& operator<<(const Doc& right); // Like above, but automatically lifts string to a Doc. Doc& operator<<(const std::string& right); // Like above, but converts right to a string first. template<typename T> Doc& operator<<(const T& right) { std::ostringstream os; os << right; return *this << os.str(); } // Indent a doc stream. friend Doc Indent(int indent, const Doc& doc); // Wadler's `layout` std::string str(); private: std::vector<DocAtom> stream_; }; // DSL functions // Render vectors of docs with a separator. e.g. PrintVec([1, 2, 3], f) -> 1f2f3 Doc PrintVec(const std::vector<Doc>& vec, const Doc& sep = Doc(", ")); // Print a constant bool value. Doc PrintBool(bool value); // Print a data type. Doc PrintDType(DataType dtype); // Print a string. Doc PrintString(const std::string& value); /*! * \brief special method to print out const scalar * \param dtype The data type * \param data The pointer to hold the data. */ template<typename T> Doc PrintConstScalar(DataType dtype, const T* data) { std::ostringstream os; if (dtype == Int(32)) { os << data[0]; } else if (dtype == Float(32)) { os << data[0] << 'f'; } else if (dtype == Bool()) { return PrintBool(data[0] != 0); } else { os << dtype << "(" << data[0] << ")"; } return Doc(os.str()); } } // namespace relay } // namespace tvm #endif // TVM_RELAY_IR_DOC_H_