vinx13 commented on a change in pull request #9306:
URL: https://github.com/apache/tvm/pull/9306#discussion_r731422912



##########
File path: src/printer/tvmscript_printer.cc
##########
@@ -1279,12 +1279,105 @@ 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());

Review comment:
       As the block syntax refactor has been merged, please rebase and check if 
the code here is up to date

##########
File path: src/tir/schedule/error.cc
##########
@@ -24,29 +24,41 @@ namespace tir {
 String ScheduleError::RenderReport(const String& primitive) const {
   IRModule mod = this->mod();
   std::ostringstream os;
-  os << "ScheduleError: An error occurred in the schedule primitive '" << 
primitive
-     << "'.\n\nThe IR is:\n"
-     << AsTVMScript(mod);
+
+  // get locations of interest
   Array<ObjectRef> locs = LocationsOfInterest();
+  std::unordered_set<ObjectRef, ObjectPtrHash, ObjectPtrEqual> loc_set;
   int n_locs = locs.size();
   std::vector<String> roi_names;
   roi_names.reserve(n_locs);
+  std::string msg = DetailRenderTemplate();
   if (n_locs > 0) {
-    os << "Regions of interest:\n";
-    for (const ObjectRef& obj : locs) {
-      String name = obj->GetTypeKey() + '#' + std::to_string(roi_names.size());
-      os << name << "\n" << obj;
+    for (int i = 0; i < n_locs; ++i) {
+      std::string name = locs[i]->GetTypeKey() + '#' + 
std::to_string(roi_names.size());
+      std::string src = "{" + std::to_string(i) + "}";
+      for (size_t pos; (pos = msg.find(src)) != std::string::npos;) {
+        msg.replace(pos, src.length(), name);
+      }
       roi_names.emplace_back(std::move(name));
-    }
-    os << "\n";
-  }
-  std::string msg = DetailRenderTemplate();
-  for (int i = 0; i < n_locs; ++i) {
-    std::string src = "{" + std::to_string(i) + "}";
-    for (size_t pos; (pos = msg.find(src)) != std::string::npos;) {
-      msg.replace(pos, src.length(), roi_names[i]);
+      loc_set.emplace(locs[i]);
     }
   }
+  msg = "Error message: " + std::move(msg);
+
+  // print IR module
+  runtime::TypedPackedFunc<std::string(Stmt)> annotate =
+      runtime::TypedPackedFunc<std::string(Stmt)>(
+          [&loc_set, &msg](const Stmt& expr) -> std::string {
+            auto search = loc_set.find(Downcast<ObjectRef>(expr));
+            if (search == loc_set.end()) return "";
+            return msg;

Review comment:
       there can be multiple regions of interest, we don't need to duplicate 
error message, only need to return name of the region of interest here, you 
might want to change `roi_names` to a map




-- 
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: commits-unsubscr...@tvm.apache.org

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


Reply via email to