[tvm] branch unity updated: [Unity][TVMScript] Update GlobalVar `checked_type_` when `emit_te` (#14367)

[ju...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

junrushao pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/unity by this push:
     new 043e4f1bd6 [Unity][TVMScript] Update GlobalVar `checked_type_` when `emit_te` (#14367)
043e4f1bd6 is described below

commit 043e4f1bd6cc85840460c4dfe867f32b712a365e
Author: Siyuan Feng <Hz...@sjtu.edu.cn>
AuthorDate: Thu Mar 23 00:02:29 2023 +0800

    [Unity][TVMScript] Update GlobalVar `checked_type_` when `emit_te` (#14367)
    
    The current GlobalVar generated by `emit_te` has empty `checked_type_`,
    which will fail at the well-formed check, saying
    
    ```
    Warning: This IR is not well formed: The checked_type_ of Expr I.GlobalVar("add") is nullptr.
    ```
    
    This PR fixes this issue and enables well-formed checks in parser tests.
---
 src/script/ir_builder/ir/ir.cc              | 14 ++++++++------
 tests/python/relax/test_tvmscript_parser.py |  3 +++
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/src/script/ir_builder/ir/ir.cc b/src/script/ir_builder/ir/ir.cc
index 6521df265e..906c453ba0 100644
--- a/src/script/ir_builder/ir/ir.cc
+++ b/src/script/ir_builder/ir/ir.cc
@@ -51,9 +51,10 @@ GlobalVar DeclFunction(const String& func_name, const BaseFunc& func_signature)
   CHECK(frame->functions.find(gv) == frame->functions.end())
       << "ValueError: function " << func_name << " has already been defined.";
   frame->global_var_map.Set(func_name, gv);
-  if (func_signature.defined()) {
-    frame->functions.Set(gv, func_signature);
-  }
+  frame->functions.Set(gv, func_signature);
+  ICHECK(func_signature->checked_type_.defined())
+      << "The checked_type_ of function signature must be defined.";
+  gv->checked_type_ = func_signature->checked_type_;
   return gv;
 }
 
@@ -64,9 +65,10 @@ void DefFunction(const String& func_name, const BaseFunc& func) {
       << "ValueError: function " << func_name << " does not exist, please declare it first.";
   const GlobalVar& gv = (*it).second;
   frame->functions.Set(gv, func);
-  if (func->checked_type_.defined()) {
-    gv->checked_type_ = func->checked_type_;
-  }
+  CHECK(func->checked_type_.defined())
+      << "The checked_type_ of function must be defined, but it is not defined for function `"
+      << func_name << "`.";
+  gv->checked_type_ = func->checked_type_;
 }
 
 void ModuleAttrs(Map<String, ObjectRef> attrs) {
diff --git a/tests/python/relax/test_tvmscript_parser.py b/tests/python/relax/test_tvmscript_parser.py
index 7ce7894960..e8c4c9de2a 100644
--- a/tests/python/relax/test_tvmscript_parser.py
+++ b/tests/python/relax/test_tvmscript_parser.py
@@ -35,6 +35,9 @@ def _check(
     test = parsed.script(show_meta=True)
     roundtrip_mod = tvm.script.from_source(test)
     tvm.ir.assert_structural_equal(parsed, roundtrip_mod)
+    if isinstance(parsed, IRModule) and isinstance(roundtrip_mod, IRModule):
+        assert relax.analysis.well_formed(parsed)
+        assert relax.analysis.well_formed(roundtrip_mod)
     if expect:
         tvm.ir.assert_structural_equal(parsed, expect)