vinx13 commented on a change in pull request #9306:
URL: https://github.com/apache/tvm/pull/9306#discussion_r732077240
##########
File path: src/printer/tvmscript_printer.cc
##########
@@ -1343,12 +1343,111 @@ Doc TVMScriptPrinter::PrintLoopStack() {
return res;
}
+/*!
+ * \brief The printer for TVMScript with diagnostic
+ * \details The printer obtain the precedence of the top-level operation when
printing each
+ * subexpression to decide whether or not parentheses is needed.
+ */
+class TVMScriptPrinterWithDiagnostic : public TVMScriptPrinter {
+ public:
+ explicit TVMScriptPrinterWithDiagnostic(const String& tir_prefix, bool
show_meta,
+
runtime::TypedPackedFunc<std::string(Stmt)> annotate)
+ : TVMScriptPrinter(tir_prefix, show_meta, annotate) {}
+
+ protected:
+ Doc VisitStmt_(const ForNode* op) override;
+ Doc VisitStmt_(const BlockRealizeNode* op) override;
+ Doc PrintAnnotation(const Stmt stmt, int length);
+};
+
+Doc TVMScriptPrinterWithDiagnostic::VisitStmt_(const ForNode* op) {
+ Doc doc;
+ var_not_in_headers_.insert(op->loop_var.get());
+ loop_var_map_[op->loop_var.get()] = GetRef<For>(op);
+ const auto* body = op->body.as<ForNode>();
+ bool simple_loop = op->kind == ForKind::kSerial && op->annotations.empty()
&& is_zero(op->min);
+ if (simple_loop) simple_loop_stack_.push_back(GetRef<For>(op));
+ // It is a loop that can be compressed, let the loops below print it out
+ if (simple_loop && body != nullptr) {
+ Doc result = Print(GetRef<For>(body));
+ TryDeallocVar(op->loop_var);
+ loop_var_map_.erase(op->loop_var.get());
+ return result;
+ }
+ // It is a loop that can not be compressed
+ bool print_above = !simple_loop_stack_.empty();
+ // print loops above if needed
+ if (print_above) {
+ doc << PrintLoopStack();
+ simple_loop_stack_.clear();
+ }
+ if (!simple_loop) {
+ // print current loop if needed
+ Doc current_loop;
+ current_loop << PrintLoop(GetRef<For>(op))
+ << PrintAnnotation(GetRef<Stmt>(op), doc.str().size());
+ current_loop << Doc::Indent(4, Doc::NewLine() << PrintBody(op->body));
+ doc << (print_above ? Doc::Indent(4, Doc::NewLine() << current_loop) :
current_loop);
+ } else {
+ doc << PrintAnnotation(GetRef<Stmt>(op), doc.str().size())
+ << Doc::Indent(4, Doc::NewLine() << PrintBody(op->body));
+ }
+ TryDeallocVar(op->loop_var);
+ loop_var_map_.erase(op->loop_var.get());
+ return doc;
+}
+
+Doc TVMScriptPrinterWithDiagnostic::VisitStmt_(const BlockRealizeNode* op) {
+ const auto* block_op = op->block.as<BlockNode>();
+ // print block name
+ Doc doc;
+ doc << "with " << tir_prefix_ << ".block(";
+ if (!block_op->name_hint.empty()) {
+ doc << Doc::StrLiteral(block_op->name_hint);
+ }
+ doc << "):";
+ // annotation
+ doc << PrintAnnotation(GetRef<Stmt>(block_op), doc.str().size());
+ // print block vars
+ Doc block_var = PrintBlockVars(op);
+ // print predicate, binding, read/write tensor region, annotations
+ Doc block_attr_doc = PrintBlockAttr(op);
+ // print body
+ Doc body = PrintBlockBody(block_op);
+ doc << Doc::Indent(4, block_var << block_attr_doc << Doc::NewLine() << body);
+ for (const auto& iter_var : block_op->iter_vars) {
+ TryDeallocVar(iter_var->var);
+ }
+ return doc;
+}
+
+Doc TVMScriptPrinterWithDiagnostic::PrintAnnotation(const Stmt stmt, int
length) {
+ Doc doc;
+ // annotation
+ if (annotate_ != nullptr) {
+ String annotated_stmt = annotate_(stmt);
+ if (!annotated_stmt.empty()) {
+ String underline = std::string(length, '^');
+ doc << Doc::NewLine() << underline << Doc::NewLine() << annotated_stmt;
Review comment:
`annotated_stmt` is also printed
https://github.com/apache/tvm/blob/main/src/printer/tvmscript_printer.cc#L433
need to check if the annotation is duplicated
--
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]
