[1] Review comment handled

python/tvm/tir/transform/transform.py

@@ -505,14 +505,24 @@ def HoistIfThenElse(variant=None):
 
     Parameters
     ----------
-    variant : str
+    variant : str, optional
         The variant of the pass.
+        variant can have any one of following values ["basic", None(Default)].
+
+        The basic variant supports basic hoisting scenarios where it exepects
+        the For & If Nodes are in place consecutively and does not involve
+        global scope variables or more advanced scenarios.
+
+        Default variant supports all hoisting scenarios,i.e., {"Basic" + "Advanced"}
+        supported with control with PassContext configs like below:
+
+            config={"tir.HoistIfThenElse": {"support_block_scope_hosting": True}}
 
     Returns
     -------
     fpass : tvm.transform.Pass
         The result pass
     """
-    if variant is None:
-        return _ffi_api.HoistIfThenElse()
-    return _ffi_api.HoistIfThenElseBasic()
+    if variant == "basic":
+        return _ffi_api.HoistIfThenElseBasic()
+    return _ffi_api.HoistIfThenElse()

src/tir/transforms/hoist_if_then_else.cc

@@ -176,9 +176,9 @@ class HoistCandidateSelector final : public StmtExprVisitor {
         StmtExprVisitor::VisitStmt_(op);
         RemoveRecord(GetRef<ObjectRef>(op));
         return;
+      } else {
+        return StmtExprVisitor::VisitStmt_(op);
       }
-
-      return StmtExprVisitor::VisitStmt_(op);
     }
     UpdateAttrVarList(op);
     StmtExprVisitor::VisitStmt_(op);
@@ -327,6 +327,7 @@ class HoistCandidateSelector final : public StmtExprVisitor {
     return false;
   }
 
+  // Ordered List maintains all ForNodes & AttrStmtNodes encountered in sequence
   std::vector<const Object*> ordered_list_;
   std::vector<const VarNode*> if_var_list_;
   std::unordered_set<const VarNode*> attr_var_list_;

tests/python/unittest/test_tir_transform_hoist_if.py

@@ -743,28 +743,4 @@ def test_hoisting_op_conv():
         tvm.testing.assert_allclose(t1.mean(), t2.mean(), atol=1, rtol=1e-1)
 
 if __name__ == "__main__":
-    test_hoist_top_for()
-    test_hoist_multi_var_if()
-    test_hoist_no_match_for()
-    test_no_else()
-    test_attr_stmt()
-    test_nested_for()
-    test_if_block()
-    test_multi_if()
-    test_no_hoisting_1()
-    test_no_hoisting_2()
-    test_no_hoisting_3()
-    test_no_hoisting_4()
-    test_no_hoisting_5()
-    test_no_hoisting_6()
-    test_no_hoisting_7()
-    test_hoisting_block_scope_1()
-    test_hoisting_block_scope_2()
-    test_hoisting_block_scope_3()
-    test_hoisting_block_scope_4()
-    test_hoisting_block_scope_5()
-    test_hoisting_block_scope_6()
-    test_hoisting_block_scope_7()
-
-    # Test with Conv Op
-    test_hoisting_op_conv()
+    pytest.main([__file__])