[Analysis] Optionally check structure info in well-formedness check (apache#321)

With the introduction of structure info in apache#314, the well-formedness check will report malformed whenever an Expr doesn’t have defined structure info.

However, when writing tests for well-formedness check and normalizer, usually we will manually construct the Exprs, which means their structure info are not defined most of the time. As a consequence, the well-formedness check will always complain “the Expr xxx doesn’t have structure info populated.” Therefore, when the checker fails to complain about the original reason of malformed, which means the checker is not working, the tests will still pass and we won’t be able to realize there is something wrong with the checker.

Thus, in this PR we add an optional flag to the well-formedness check. In well-formedness tests, we will turn off the structure info check so that the original reason of being malformed will be revealed correctly.

---

This PR also cleans up the DiagnosticContext parameter in the WellFormed API - the diag_ctx has been unused since the merge of apache#99.

Author: MasterJH5574

include/tvm/relax/analysis.h

@@ -200,8 +200,8 @@ TVM_DLL StructInfo EraseToWellDefined(const StructInfo& info, Map<tir::Var, Prim
  *
  * For a given pair of lhs_struct_info, rhs_struct_info. We adopt
  * the following terminology:
- * - LSet = {value | value mactches lhs_struct_info}
- * - RSet = {value | value mactches rhs_struct_info}
+ * - LSet = {value | value matches lhs_struct_info}
+ * - RSet = {value | value matches rhs_struct_info}
  *
  * See the definition of each level below.
  */
@@ -283,11 +283,14 @@ TVM_DLL StructInfo StructInfoLCA(const StructInfo& lhs, const StructInfo& rhs,
  * \brief Check if the IRModule is well formed.
  *
  * \param m the IRModule to check.
- * \param diag_ctx the diagnostic context.
+ * \param check_struct_info A boolean flag indicating if the property "every Expr
+ * must have defined structure info" will be checked.
  * \return true if the IRModule is well formed, false if not.
+ * \note By default the structure info is always checked. It is only in test cases
+ * where `check_struct_info` might be false, so that other well-formed requirements
+ * will be well tested and will not be blocked by not having structure info.
  */
-TVM_DLL bool WellFormed(const IRModule& m,
-                        Optional<DiagnosticContext> diag_ctx = Optional<DiagnosticContext>());
+TVM_DLL bool WellFormed(IRModule m, bool check_struct_info = true);
 
 /*!
  * \brief Annotate Op Pattern Kind for PrimFunc, which is used in relax FuseOps.
@@ -346,7 +349,7 @@ TVM_DLL tvm::Array<Var> FreeVars(const Expr& expr);
 TVM_DLL tvm::Array<Var> AllVars(const Expr& expr);
 
 /*!
- * \brief Get all glabal variables used in calls in expression expr.
+ * \brief Get all global variables used in calls in expression expr.
  *
  * \param expr the expression.
  *
@@ -355,7 +358,7 @@ TVM_DLL tvm::Array<Var> AllVars(const Expr& expr);
 TVM_DLL tvm::Array<GlobalVar> CalledGlobalVars(const Expr& expr);
 
 /*!
- * \brief Get all glabal variables from expression expr.
+ * \brief Get all global variables from expression expr.
  *
  * AllVars is a superset of BoundVars and FreeVars.
  * The union of BoundVars and FreeVars is Allvars.
@@ -402,15 +405,15 @@ TVM_DLL Map<String, Array<Binding>> NameToBinding(const Function& fn);
  * \brief Get the use-def chain of variables inside a dataflow block.
  *
  * \param dfb The dataflow block to be analyzed.
- * \return A map mapping variable definitoins to a set of uses.
+ * \return A map mapping variable definitions to a set of uses.
  */
 TVM_DLL Map<Var, Array<Var>> DataflowBlockUseDef(const DataflowBlock& dfb);
 
 /*!
  * \brief Get the use-def chain of variables inside a function.
  *
  * \param fn The function to be analyzed.
- * \return A map from variable definitoins to a set of uses and variables needed by return value.
+ * \return A map from variable definitions to a set of uses and variables needed by return value.
  */
 std::pair<Map<Var, Array<Var>>, Array<Var>> FunctionUseDef(const Function& fn);
 

python/tvm/relax/analysis/analysis.py

@@ -197,20 +197,30 @@ def post_order_visit(expr, fvisit):
     return _ffi_api.post_order_visit(expr, fvisit)  # type: ignore
 
 
-def well_formed(mod: tvm.IRModule) -> bool:
+def well_formed(mod: tvm.IRModule, check_struct_info: bool = True) -> bool:
     """Check if the IRModule is well formed.
 
     Parameters
     ----------
     mod : tvm.IRModule
         The input IRModule.
 
+    check_struct_info : bool
+        A boolean flag indicating if the property "every Expr must
+        have defined structure info" will be checked.
+
     Returns
     -------
     ret: bool
         True if the IRModule is well formed, False if not.
+
+    Note
+    ----
+    By default the structure info is always checked. It is only in test cases
+    where `check_struct_info` might be false, so that other well-formed requirements
+    will be well tested and will not be blocked by not having structure info.
     """
-    return _ffi_api.well_formed(mod)  # type: ignore
+    return _ffi_api.well_formed(mod, check_struct_info)  # type: ignore
 
 
 def get_var2val(func: Function) -> Dict[Var, Expr]: