Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 1 commit into from
Feb 7, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 12 additions & 3 deletions include/tvm/node/script_printer.h
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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";
Expand Down
39 changes: 39 additions & 0 deletions include/tvm/script/printer/ir_docsifier.h
Original file line number Diff line number Diff line change
Expand Up @@ -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));
Expand Down
2 changes: 2 additions & 0 deletions python/tvm/runtime/object_path.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
57 changes: 45 additions & 12 deletions python/tvm/runtime/script_printer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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,
Expand All @@ -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
Expand All @@ -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,
},
)

Expand All @@ -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

Expand All @@ -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
-------
Expand All @@ -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,
),
)

Expand All @@ -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.

Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down
4 changes: 2 additions & 2 deletions python/tvm/script/printer/doc_printer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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.

Expand Down
13 changes: 12 additions & 1 deletion src/node/script_printer.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
12 changes: 6 additions & 6 deletions src/node/structural_equal.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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.
Expand Down
Loading