[TIR] Expose UndefinedVars to Python

python/tvm/tir/analysis/analysis.py

@@ -16,7 +16,7 @@
 # under the License.
 """Wrapping existing analysis utils."""
 # pylint: disable=invalid-name
-from typing import Dict, List, Union
+from typing import Dict, List, Optional, Union
 
 import tvm
 from tvm import Object
@@ -211,7 +211,7 @@ def calculate_allocated_bytes(
     ----------
     func_or_mod: Union[PrimFunc, IRModule]
         The function or module to be detected. If a module is passed, allocated
-        memory is calcualted for all PrimFuncs inside the module
+        memory is calculated for all PrimFuncs inside the module
 
     Returns
     -------
@@ -266,6 +266,26 @@ def estimate_tir_flops(stmt_or_mod: Union[Stmt, IRModule]) -> float:
 # introduce a cycling dependency. We make do with Object.
 
 
+def undefined_vars(node: Union[Stmt, PrimExpr], defs: Optional[List[Var]] = None) -> List[Var]:
+    """Find undefined vars in a TIR statement or expression.
+
+    Parameters
+    ----------
+    node: Union[Stmt, PrimExpr]
+        The TIR statement or expression to be checked.
+
+    defs: Optional[List[Var]]
+        The vars that is defined
+
+    Returns
+    -------
+    result : List[Var]
+        The undefined vars.
+    """
+    defs = defs or []
+    return _ffi_api.UndefinedVars(node, defs)  # type: ignore # pylint: disable=no-member
+
+
 def get_prim_func_arg_and_result_memory_constraints(
     func: PrimFunc, relay_func_type: Object
 ) -> List[str]:
@@ -388,7 +408,7 @@ def find_anchor_block(mod: IRModule) -> Block:
 
 
 def get_vtcm_compaction_passes() -> List[tvm.transform.Pass]:
-    """Utility function to get the list of lowering passes to be applied to calculate thecompacted
+    """Utility function to get the list of lowering passes to be applied to calculate the compacted
     VTCM allocation size
 
     Returns

src/tir/analysis/var_use_def_analysis.cc

@@ -199,5 +199,14 @@ Array<Var> UndefinedVars(const PrimExpr& expr, const Array<Var>& args) {
   return m.undefined_;
 }
 
+TVM_REGISTER_GLOBAL("tir.analysis.UndefinedVars").set_body([](TVMArgs args, TVMRetValue* rv) {
+  if (args.size() == 2 && args[0].IsObjectRef<Stmt>()) {
+    *rv = UndefinedVars(args[0].AsObjectRef<Stmt>(), args[1]);
+  } else if (args.size() == 2 && args[0].IsObjectRef<PrimExpr>()) {
+    *rv = UndefinedVars(args[0].AsObjectRef<PrimExpr>(), args[1]);
+  } else {
+    LOG(FATAL) << "either UndefinedVars(stmt, args) or UndefinedVars(expr, args) is expected";
+  }
+});
 }  // namespace tir
 }  // namespace tvm