[TVMScript] Schedule error reporting with new TVMScript printer (#13921)

This PR makes the following changes over TVMScript printer:

1. Apply the new TVMScript printer to report and render the schedule error, and keeps effect same as old `AsTVMScriptWithDiagnostic`.
2. Introduce more annotating and underlining interfaces for TVMScript ptrinter
    - `obj_to_annotate`: add comments to the final `StmtDoc` from `ObjectRef`
    - `obj_to_underline`: underline the final `Doc` from `ObjectRef`
    - `path_to_annotate`: add comments to the final `StmtDoc` from `ObjectPath`
3. Beautify the underline logic by introducing the `underline_exempted`, to exempt some useless underlines, e.g. indent, comments and doc string.

demo:
1. `*_to_annotate`: for function:
```python
@T.prim_func
def _func():
  T.evaluate(0)
  T.evaluate(1)
  T.evaluate(2)
  T.evaluate(3)
  T.evaluate(4)
  T.evaluate(5)
  T.evaluate(6)
  T.evaluate(7)
```
both the results of 
```python
_func.script(
  path_to_annotate={
    ObjectPath.root().attr("body").attr("seq").array_index(1): "annotation 1",
    ObjectPath.root().attr("body").attr("seq").array_index(3): "annotation 3",
    ObjectPath.root().attr("body").attr("seq").array_index(5): "annotation 5",
    ObjectPath.root().attr("body").attr("seq").array_index(7): "annotation 7",
  }
)
```
and
```python
_func.script(
  obj_to_annotate={
    _func.body.seq[1]: "annotation 1",
    _func.body.seq[3]: "annotation 3",
    _func.body.seq[5]: "annotation 5",
    _func.body.seq[7]: "annotation 7",
  }
)
```
are
```python
# from tvm.script import tir as T
@T.prim_func
def main():
  T.evaluate(0)
  T.evaluate(1)  # annotation 1
  T.evaluate(2)
  T.evaluate(3)  # annotation 3
  T.evaluate(4)
  T.evaluate(5)  # annotation 5
  T.evaluate(6)
  T.evaluate(7)  # annotation 7
```
2. `obj_to_underline`: for function
```python
@T.prim_func
def func(a: T.int32, b: T.int32):
  T.evaluate(a)
  T.evaluate(b)
  T.evaluate(a)
  T.evaluate(b)
  T.evaluate(a)
  T.evaluate(b)
```
the result of `func.script(obj_to_underline=[func.params[0]])` is
```python
# from tvm.script import tir as T
@T.prim_func
def main(a: T.int32, b: T.int32):
  T.evaluate(a)
             ^
  T.evaluate(b)
  T.evaluate(a)
             ^
  T.evaluate(b)
  T.evaluate(a)
             ^
  T.evaluate(b)
```

include/tvm/node/script_printer.h

@@ -59,10 +59,16 @@ class PrinterConfigNode : public Object {
   bool print_line_numbers = false;
   /*! \brief Number of context lines to print around the underlined text */
   int num_context_lines = -1;
-  /*! \brief Object path to be underlined */
-  Optional<ObjectPath> path_to_underline = NullOpt;
   /*! \brief Whether to output with syntax sugar, set false for complete printing. */
   bool syntax_sugar = true;
+  /* \brief Object path to be underlined */
+  Array<ObjectPath> path_to_underline = Array<ObjectPath>();
+  /*! \brief Object path to be annotated. */
+  Map<ObjectPath, String> path_to_annotate = Map<ObjectPath, String>();
+  /*! \brief Object to be underlined. */
+  Array<ObjectRef> obj_to_underline = Array<ObjectRef>();
+  /*! \brief Object to be annotated. */
+  Map<ObjectRef, String> obj_to_annotate = Map<ObjectRef, String>();
 
   void VisitAttrs(AttrVisitor* v) {
     v->Visit("ir_prefix", &ir_prefix);
@@ -73,8 +79,11 @@ class PrinterConfigNode : public Object {
     v->Visit("indent_spaces", &indent_spaces);
     v->Visit("print_line_numbers", &print_line_numbers);
     v->Visit("num_context_lines", &num_context_lines);
-    v->Visit("path_to_underline", &path_to_underline);
     v->Visit("syntax_sugar", &syntax_sugar);
+    v->Visit("path_to_underline", &path_to_underline);
+    v->Visit("path_to_annotate", &path_to_annotate);
+    v->Visit("obj_to_underline", &obj_to_underline);
+    v->Visit("obj_to_annotate", &obj_to_annotate);
   }
 
   static constexpr const char* _type_key = "node.PrinterConfig";

include/tvm/script/printer/ir_docsifier.h

@@ -263,11 +263,50 @@ inline void FrameNode::ExitWithScope() {
   }
 }
 
+template <class TDoc>
+inline static void AddDocDecoration(const Doc& d, const ObjectRef& obj, const ObjectPath& path,
+                                    const PrinterConfig& cfg) {
+  if (cfg->obj_to_annotate.count(obj)) {
+    if (const auto* stmt = d.as<StmtDocNode>()) {
+      if (stmt->comment.defined()) {
+        stmt->comment = stmt->comment.value() + "\n" + cfg->obj_to_annotate.at(obj);
+      } else {
+        stmt->comment = cfg->obj_to_annotate.at(obj);
+      }
+    } else {
+      LOG(WARNING) << "Expect StmtDoc to be annotated for object " << obj << ", but got "
+                   << Downcast<TDoc>(d)->_type_key;
+    }
+  }
+  for (const ObjectRef& o : cfg->obj_to_underline) {
+    if (o.same_as(obj)) {
+      cfg->path_to_underline.push_back(path);
+    }
+  }
+  for (const auto& pair : cfg->path_to_annotate) {
+    ObjectPath p = pair.first;
+    String attn = pair.second;
+    if (p->IsPrefixOf(path) && path->IsPrefixOf(p)) {
+      if (const auto* stmt = d.as<StmtDocNode>()) {
+        if (stmt->comment.defined()) {
+          stmt->comment = stmt->comment.value() + "\n" + attn;
+        } else {
+          stmt->comment = attn;
+        }
+      } else {
+        LOG(WARNING) << "Expect StmtDoc to be annotated at object path " << p << ", but got "
+                     << Downcast<TDoc>(d)->_type_key;
+      }
+    }
+  }
+}
+
 template <class TDoc>
 inline TDoc IRDocsifierNode::AsDoc(const ObjectRef& obj, const ObjectPath& path) const {
   if (obj.defined()) {
     Doc d = IRDocsifier::vtable()(dispatch_tokens.back(), obj, path, GetRef<IRDocsifier>(this));
     d->source_paths.push_back(path);
+    AddDocDecoration<TDoc>(d, obj, path, cfg);
     return Downcast<TDoc>(d);
   }
   return Downcast<TDoc>(LiteralDoc::None(path));

python/tvm/runtime/object_path.py

@@ -89,6 +89,8 @@ def map_value(self, key) -> "ObjectPath":
     def missing_map_entry(self) -> "ObjectPath":
         return _ffi_node_api.ObjectPathMissingMapEntry(self)
 
+    __hash__ = Object.__hash__
+
 
 @tvm._ffi.register_object("RootPath")
 class RootPath(ObjectPath):

python/tvm/runtime/script_printer.py

@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """Configuration of TVMScript printer"""
-from typing import Optional
+from typing import List, Dict, Optional
 
 from tvm._ffi import register_object
 from tvm.runtime import Object
@@ -38,8 +38,11 @@ class PrinterConfig(Object):
     indent_spaces: int
     print_line_numbers: bool
     num_context_lines: int
-    path_to_underline: Optional[ObjectPath]
     syntax_sugar: bool
+    path_to_underline: Optional[List[ObjectPath]]
+    path_to_annotate: Optional[Dict[ObjectPath, str]]
+    obj_to_underline: Optional[List[Object]]
+    obj_to_annotate: Optional[Dict[Object, str]]
 
     def __init__(
         self,
@@ -54,8 +57,11 @@ def __init__(
         indent_spaces: int = 4,
         print_line_numbers: bool = False,
         num_context_lines: Optional[int] = None,
-        path_to_underline: Optional[ObjectPath] = None,
         syntax_sugar: bool = True,
+        path_to_underline: Optional[List[ObjectPath]] = None,
+        path_to_annotate: Optional[Dict[ObjectPath, str]] = None,
+        obj_to_underline: Optional[List[Object]] = None,
+        obj_to_annotate: Optional[Dict[Object, str]] = None,
     ) -> None:
         if num_context_lines is None:
             num_context_lines = -1
@@ -72,8 +78,11 @@ def __init__(
                 "indent_spaces": indent_spaces,
                 "print_line_numbers": print_line_numbers,
                 "num_context_lines": num_context_lines,
-                "path_to_underline": path_to_underline,
                 "syntax_sugar": syntax_sugar,
+                "path_to_underline": path_to_underline,
+                "path_to_annotate": path_to_annotate,
+                "obj_to_underline": obj_to_underline,
+                "obj_to_annotate": obj_to_annotate,
             },
         )
 
@@ -98,8 +107,11 @@ def script(
         indent_spaces: int = 4,
         print_line_numbers: bool = False,
         num_context_lines: int = -1,
-        path_to_underline: Optional[ObjectPath] = None,
         syntax_sugar: bool = True,
+        path_to_underline: Optional[List[ObjectPath]] = None,
+        path_to_annotate: Optional[Dict[ObjectPath, str]] = None,
+        obj_to_underline: Optional[List[Object]] = None,
+        obj_to_annotate: Optional[Dict[Object, str]] = None,
     ) -> str:
         """Print TVM IR into TVMScript text format
 
@@ -125,10 +137,16 @@ def script(
             Whether to print line numbers
         num_context_lines : int = -1
             The number of lines of context to print before and after the line to underline.
-        path_to_underline : Optional[ObjectPath] = None
-            Object path to be underlined
         syntax_sugar: bool = True
              Whether to output with syntax sugar, set false for complete printing.
+        path_to_underline : Optional[List[ObjectPath]] = None
+            Object path to be underlined
+        path_to_annotate : Optional[Dict[ObjectPath, str]] = None
+            Object path to be annotated
+        obj_to_underline : Optional[List[Object]] = None
+            Object to be underlined
+        obj_to_annotate : Optional[Dict[Object, str]] = None
+            Object to be annotated
 
         Returns
         -------
@@ -148,8 +166,11 @@ def script(
                 indent_spaces=indent_spaces,
                 print_line_numbers=print_line_numbers,
                 num_context_lines=num_context_lines,
-                path_to_underline=path_to_underline,
                 syntax_sugar=syntax_sugar,
+                path_to_underline=path_to_underline,
+                path_to_annotate=path_to_annotate,
+                obj_to_underline=obj_to_underline,
+                obj_to_annotate=obj_to_annotate,
             ),
         )
 
@@ -168,8 +189,11 @@ def show(
         indent_spaces: int = 4,
         print_line_numbers: bool = False,
         num_context_lines: int = -1,
-        path_to_underline: Optional[ObjectPath] = None,
         syntax_sugar: bool = True,
+        path_to_underline: Optional[List[ObjectPath]] = None,
+        path_to_annotate: Optional[Dict[ObjectPath, str]] = None,
+        obj_to_underline: Optional[List[Object]] = None,
+        obj_to_annotate: Optional[Dict[Object, str]] = None,
     ) -> None:
         """A sugar for print highlighted TVM script.
 
@@ -200,10 +224,16 @@ def show(
             Whether to print line numbers
         num_context_lines : int = -1
             The number of lines of context to print before and after the line to underline.
-        path_to_underline : Optional[ObjectPath] = None
-            Object path to be underlined
         syntax_sugar: bool = True
              Whether to output with syntax sugar, set false for complete printing.
+        path_to_underline : Optional[List[ObjectPath]] = None
+            Object path to be underlined
+        path_to_annotate : Optional[Dict[ObjectPath, str]] = None
+            Object path to be annotated
+        obj_to_underline : Optional[List[Object]] = None
+            Object to be underlined
+        obj_to_annotate : Optional[Dict[Object, str]] = None
+            Object to be annotated
         """
         from tvm.script.highlight import (  # pylint: disable=import-outside-toplevel
             cprint,
@@ -221,8 +251,11 @@ def show(
                 indent_spaces=indent_spaces,
                 print_line_numbers=print_line_numbers,
                 num_context_lines=num_context_lines,
-                path_to_underline=path_to_underline,
                 syntax_sugar=syntax_sugar,
+                path_to_underline=path_to_underline,
+                path_to_annotate=path_to_annotate,
+                obj_to_underline=obj_to_underline,
+                obj_to_annotate=obj_to_annotate,
             ),
             style=style,
             black_format=black_format,

python/tvm/script/printer/doc_printer.py

@@ -16,7 +16,7 @@
 # under the License.
 """Functions to print doc into text format"""
 
-from typing import Optional
+from typing import List, Optional
 
 from tvm.runtime import ObjectPath
 from tvm.runtime.script_printer import PrinterConfig
@@ -30,7 +30,7 @@ def to_python_script(
     indent_spaces: int = 4,
     print_line_numbers: bool = False,
     num_context_lines: Optional[int] = None,
-    path_to_underline: Optional[ObjectPath] = None,
+    path_to_underline: Optional[List[ObjectPath]] = None,
 ) -> str:
     """Convert Doc into Python script.
 

src/node/script_printer.cc

@@ -68,7 +68,18 @@ PrinterConfig::PrinterConfig(Map<String, ObjectRef> config_dict) {
     n->num_context_lines = Downcast<IntImm>(v)->value;
   }
   if (auto v = config_dict.Get("path_to_underline")) {
-    n->path_to_underline = Downcast<ObjectPath>(v);
+    n->path_to_underline = Downcast<Optional<Array<ObjectPath>>>(v).value_or(Array<ObjectPath>());
+  }
+  if (auto v = config_dict.Get("path_to_annotate")) {
+    n->path_to_annotate =
+        Downcast<Optional<Map<ObjectPath, String>>>(v).value_or(Map<ObjectPath, String>());
+  }
+  if (auto v = config_dict.Get("obj_to_underline")) {
+    n->obj_to_underline = Downcast<Optional<Array<ObjectRef>>>(v).value_or(Array<ObjectRef>());
+  }
+  if (auto v = config_dict.Get("obj_to_annotate")) {
+    n->obj_to_annotate =
+        Downcast<Optional<Map<ObjectRef, String>>>(v).value_or(Map<ObjectRef, String>());
   }
   if (auto v = config_dict.Get("syntax_sugar")) {
     n->syntax_sugar = Downcast<IntImm>(v)->value;

src/node/structural_equal.cc

@@ -326,9 +326,9 @@ class SEqualHandlerDefault::Impl {
       if (first_mismatch_->defined()) {
         oss << " at " << first_mismatch_->value()->lhs_path;
         if (root_lhs_.defined()) {
-          Map<String, ObjectRef> dict = {{"path_to_underline", first_mismatch_->value()->lhs_path},
-                                         {"syntax_sugar", Bool(false)}};
-          PrinterConfig cfg(dict);
+          PrinterConfig cfg;
+          cfg->syntax_sugar = false;
+          cfg->path_to_underline.push_back(first_mismatch_->value()->lhs_path);
           // The TVMScriptPrinter::Script will fallback to Repr printer,
           // if the root node to print is not supported yet,
           // e.g. Relay nodes, ArrayNode, MapNode, etc.
@@ -341,9 +341,9 @@ class SEqualHandlerDefault::Impl {
       if (first_mismatch_->defined()) {
         oss << " at " << first_mismatch_->value()->rhs_path;
         if (root_rhs_.defined()) {
-          Map<String, ObjectRef> dict = {{"path_to_underline", first_mismatch_->value()->rhs_path},
-                                         {"syntax_sugar", Bool(false)}};
-          PrinterConfig cfg(dict);
+          PrinterConfig cfg;
+          cfg->syntax_sugar = false;
+          cfg->path_to_underline.push_back(first_mismatch_->value()->rhs_path);
           // The TVMScriptPrinter::Script will fallback to Repr printer,
           // if the root node to print is not supported yet,
           // e.g. Relay nodes, ArrayNode, MapNode, etc.