From 3827ac13ed14cfa086f04293a172de1bdcaaacc4 Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Tue, 15 Aug 2023 14:14:11 +0000 Subject: [PATCH 1/2] [TVMScript] Optionally output the address as part of variable names When debugging IRModule transformations, it can be useful to know the exact C++ address of an object. For example, to determine whether an undefined TIR variable was caused by erroneous insertion of a new variable, or from failing to remove a previous variable. While TVMScript does de-duplicate all names within a single print statement, there isn't a convenient way to identify which variable in TVMScript corresponds to a specific variable in C++ logging statements. This commit adds `show_object_address` to the `PrinterConfig`. If false (the default), no change is made to the TVMScript variable names. If true, the address of the C++ object is appended to the variable name. For example, printing a `tir.Var('my_name','int64')` as `"my_name_0x1234abcd"`. --- include/tvm/node/script_printer.h | 3 ++ python/tvm/runtime/script_printer.py | 11 +++++ src/node/script_printer.cc | 4 ++ src/script/printer/ir_docsifier.cc | 10 ++++- .../unittest/test_tvmscript_printer_tir.py | 41 +++++++++++++++++++ 5 files changed, 68 insertions(+), 1 deletion(-) diff --git a/include/tvm/node/script_printer.h b/include/tvm/node/script_printer.h index c65394f7b7e1..f25ffa1caa1c 100644 --- a/include/tvm/node/script_printer.h +++ b/include/tvm/node/script_printer.h @@ -68,6 +68,8 @@ class PrinterConfigNode : public Object { int num_context_lines = -1; /*! \brief Whether to output with syntax sugar, set false for complete printing. */ bool syntax_sugar = true; + /*! \brief Whether variable names should include the object's address */ + bool show_object_address = false; /* \brief Object path to be underlined */ Array path_to_underline = Array(); /*! \brief Object path to be annotated. */ @@ -91,6 +93,7 @@ class PrinterConfigNode : public Object { v->Visit("print_line_numbers", &print_line_numbers); v->Visit("num_context_lines", &num_context_lines); v->Visit("syntax_sugar", &syntax_sugar); + v->Visit("show_object_address", &show_object_address); v->Visit("path_to_underline", &path_to_underline); v->Visit("path_to_annotate", &path_to_annotate); v->Visit("obj_to_underline", &obj_to_underline); diff --git a/python/tvm/runtime/script_printer.py b/python/tvm/runtime/script_printer.py index 269cab8e5d4d..2ed2b8ddd4bc 100644 --- a/python/tvm/runtime/script_printer.py +++ b/python/tvm/runtime/script_printer.py @@ -40,6 +40,7 @@ class PrinterConfig(Object): print_line_numbers: bool num_context_lines: int syntax_sugar: bool + show_object_address: bool path_to_underline: Optional[List[ObjectPath]] path_to_annotate: Optional[Dict[ObjectPath, str]] obj_to_underline: Optional[List[Object]] @@ -60,6 +61,7 @@ def __init__( print_line_numbers: bool = False, num_context_lines: Optional[int] = None, syntax_sugar: bool = True, + show_object_address: bool = True, path_to_underline: Optional[List[ObjectPath]] = None, path_to_annotate: Optional[Dict[ObjectPath, str]] = None, obj_to_underline: Optional[List[Object]] = None, @@ -79,6 +81,7 @@ def __init__( "print_line_numbers": print_line_numbers, "num_context_lines": num_context_lines, "syntax_sugar": syntax_sugar, + "show_object_address": show_object_address, "path_to_underline": path_to_underline, "path_to_annotate": path_to_annotate, "obj_to_underline": obj_to_underline, @@ -119,6 +122,7 @@ def script( print_line_numbers: bool = False, num_context_lines: int = -1, syntax_sugar: bool = True, + show_object_address: bool = False, path_to_underline: Optional[List[ObjectPath]] = None, path_to_annotate: Optional[Dict[ObjectPath, str]] = None, obj_to_underline: Optional[List[Object]] = None, @@ -153,6 +157,8 @@ def script( The number of lines of context to print before and after the line to underline. syntax_sugar: bool = True Whether to output with syntax sugar, set false for complete printing. + show_object_address: bool = False + Whether to include the object's adddress as part of the TVMScript name path_to_underline : Optional[List[ObjectPath]] = None Object path to be underlined path_to_annotate : Optional[Dict[ObjectPath, str]] = None @@ -182,6 +188,7 @@ def script( print_line_numbers=print_line_numbers, num_context_lines=num_context_lines, syntax_sugar=syntax_sugar, + show_object_address=show_object_address, path_to_underline=path_to_underline, path_to_annotate=path_to_annotate, obj_to_underline=obj_to_underline, @@ -206,6 +213,7 @@ def show( print_line_numbers: bool = False, num_context_lines: int = -1, syntax_sugar: bool = True, + show_object_address: bool = True, path_to_underline: Optional[List[ObjectPath]] = None, path_to_annotate: Optional[Dict[ObjectPath, str]] = None, obj_to_underline: Optional[List[Object]] = None, @@ -245,6 +253,8 @@ def show( The number of lines of context to print before and after the line to underline. syntax_sugar: bool = True Whether to output with syntax sugar, set false for complete printing. + show_object_address: bool = False + Whether to include the object's adddress as part of the TVMScript name path_to_underline : Optional[List[ObjectPath]] = None Object path to be underlined path_to_annotate : Optional[Dict[ObjectPath, str]] = None @@ -272,6 +282,7 @@ def show( print_line_numbers=print_line_numbers, num_context_lines=num_context_lines, syntax_sugar=syntax_sugar, + show_object_address=show_object_address, path_to_underline=path_to_underline, path_to_annotate=path_to_annotate, obj_to_underline=obj_to_underline, diff --git a/src/node/script_printer.cc b/src/node/script_printer.cc index 8293af402ed9..28e72be78945 100644 --- a/src/node/script_printer.cc +++ b/src/node/script_printer.cc @@ -88,6 +88,10 @@ PrinterConfig::PrinterConfig(Map config_dict) { if (auto v = config_dict.Get("syntax_sugar")) { n->syntax_sugar = Downcast(v)->value; } + if (auto v = config_dict.Get("show_object_address")) { + n->show_object_address = Downcast(v)->value; + } + this->data_ = std::move(n); } diff --git a/src/script/printer/ir_docsifier.cc b/src/script/printer/ir_docsifier.cc index 7cd27057f4d4..7dc971bd2c35 100644 --- a/src/script/printer/ir_docsifier.cc +++ b/src/script/printer/ir_docsifier.cc @@ -21,6 +21,8 @@ #include #include +#include + #include "./utils.h" namespace tvm { @@ -29,7 +31,13 @@ namespace printer { IdDoc IRDocsifierNode::Define(const ObjectRef& obj, const Frame& frame, const String& name_hint) { ICHECK(obj2info.find(obj) == obj2info.end()) << "Duplicated object: " << obj; - String name = GenerateUniqueName(name_hint, this->defined_names); + String name = name_hint; + if (cfg->show_object_address) { + std::stringstream stream; + stream << name << "_" << obj.get(); + name = stream.str(); + } + name = GenerateUniqueName(name, this->defined_names); this->defined_names.insert(name); DocCreator doc_factory = [name]() { return IdDoc(name); }; obj2info.insert({obj, VariableInfo{std::move(doc_factory), name}}); diff --git a/tests/python/unittest/test_tvmscript_printer_tir.py b/tests/python/unittest/test_tvmscript_printer_tir.py index e6334553d64f..58fe0bb45ad7 100644 --- a/tests/python/unittest/test_tvmscript_printer_tir.py +++ b/tests/python/unittest/test_tvmscript_printer_tir.py @@ -15,6 +15,9 @@ # specific language governing permissions and limitations # under the License. # pylint: disable=missing-docstring + +import re + import tvm.testing from tvm import ir, tir from tvm.ir import Range @@ -798,5 +801,43 @@ def main(): _assert_print(root_block_explicitly, expected_output) +def test_variable_with_cpp_address(): + """The show_object_address option displays the C++ addressess + + Because the C++ address may vary with each execution, the output + produced with this option cannot be compared to a fixed string. + Instead, this test uses the normal script output to generate a + regular expression against with the test output must match. The + regular expression validates that all names have been appended + with "_0x" followed by a hexadecimal number, and that the address + is the same for each variable. + """ + from tvm.script import tir as T + + # The test function has all named objects suffixed with "_name", + # to avoid spurious replacement when generating the expected + # regex. + @T.prim_func(private=True) + def func(a_name: T.handle): + N_name = T.int64() + A_name = T.match_buffer(a_name, N_name, "float32") + for i_name in range(N_name): + A_name[i_name] = A_name[i_name] + 1.0 + + without_address = func.script(show_object_address=False) + script = func.script(show_object_address=True) + + expected_regex = re.escape(without_address) + for name in ["main", "a_name", "A_name", "N_name", "i_name"]: + # Replace all occurrences with a backref to an earlier match + expected_regex = expected_regex.replace(name, rf"(?P={name})") + # Then replace the first such backref with a capturing group. + expected_regex = expected_regex.replace( + rf"(?P={name})", rf"(?P<{name}>{name}_0x[A-Fa-f0-9]+)", 1 + ) + + assert re.match(expected_regex, script) + + if __name__ == "__main__": tvm.testing.main() From 9d98bca136596f28356f2800fdbe80da2d7c9441 Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Sat, 19 Aug 2023 09:17:13 -0500 Subject: [PATCH 2/2] Updated unit test to target TVM main The `private=True` branch is only on the unity branch. --- tests/python/unittest/test_tvmscript_printer_tir.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/python/unittest/test_tvmscript_printer_tir.py b/tests/python/unittest/test_tvmscript_printer_tir.py index 58fe0bb45ad7..572c22e80e51 100644 --- a/tests/python/unittest/test_tvmscript_printer_tir.py +++ b/tests/python/unittest/test_tvmscript_printer_tir.py @@ -817,7 +817,7 @@ def test_variable_with_cpp_address(): # The test function has all named objects suffixed with "_name", # to avoid spurious replacement when generating the expected # regex. - @T.prim_func(private=True) + @T.prim_func def func(a_name: T.handle): N_name = T.int64() A_name = T.match_buffer(a_name, N_name, "float32") @@ -828,7 +828,7 @@ def func(a_name: T.handle): script = func.script(show_object_address=True) expected_regex = re.escape(without_address) - for name in ["main", "a_name", "A_name", "N_name", "i_name"]: + for name in ["a_name", "A_name", "N_name", "i_name"]: # Replace all occurrences with a backref to an earlier match expected_regex = expected_regex.replace(name, rf"(?P={name})") # Then replace the first such backref with a capturing group.