diff --git a/include/tvm/tirx/analysis.h b/include/tvm/tirx/analysis.h
index 6323c03f96e6..eb90ca5226c2 100644
--- a/include/tvm/tirx/analysis.h
+++ b/include/tvm/tirx/analysis.h
@@ -106,22 +106,6 @@ TVM_DLL ffi::Array<Var> UndefinedVars(const PrimExpr& expr, const ffi::Array<Var
  */
 TVM_DLL CallEffectKind SideEffect(const PrimExpr& expr);
 
-/*!
- * \brief Whether the given Stmt uses any var in the given variable set.
- * \param stmt The Stmt to be checked.
- * \param vset_contains The check function to see if a var is in the variable set.
- * \return Whether `stmt` uses any var in the given variable set.
- */
-TVM_DLL bool UsesVar(const Stmt& stmt, std::function<bool(const VarNode*)> vset_contains);
-
-/*!
- * \brief Whether the given PrimExpr uses any var in the given variable set.
- * \param expr The PrimExpr to be checked.
- * \param vset_contains The check function to see if var is in the variable set.
- * \return Whether `expr` uses any var in the given variable set.
- */
-TVM_DLL bool UsesVar(const PrimExpr& expr, std::function<bool(const VarNode*)> vset_contains);
-
 /*!
  * \brief Verifies whether the IR stmt or Expr is in SSA form.
  *  That is: each Var is defined and assigned once(in Let/For)
diff --git a/src/arith/detect_linear_equation.cc b/src/arith/detect_linear_equation.cc
index f00bb889e638..a2a8ffdd5df2 100644
--- a/src/arith/detect_linear_equation.cc
+++ b/src/arith/detect_linear_equation.cc
@@ -110,7 +110,11 @@ class LinearEqDetector : public ExprFunctor<LinearEqEntry(const Expr&, const Pri
   }
   LinearEqEntry VisitExprDefault_(const ffi::Object* op, const PrimExpr& e) final {
     if (fail_) return LinearEqEntry();
-    if (UsesVar(e, [this](const VarNode* var) { return var == var_.get(); })) {
+    auto walkfn = [this](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+      return var.get() == var_.get() ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                     : ffi::WalkResult::Advance();
+    };
+    if (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(e, walkfn).has_value()) {
       fail_ = true;
       return LinearEqEntry();
     } else {
@@ -157,11 +161,15 @@ ffi::Array<PrimExpr> DetectLinearEquation(const PrimExpr& e, const ffi::Array<Pr
 
   std::unordered_set<const VarNode*> vset;
   auto vset_contains = [&](const VarNode* node) { return vset.count(node) != 0; };
+  auto walkfn = [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+    return vset_contains(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                    : ffi::WalkResult::Advance();
+  };
 
   for (size_t i = vars.size(); i > 1; --i) {
     vset.insert(vars[i - 1].get());
     // The previous coeff contains the variable
-    if (UsesVar(coeff[i - 2], vset_contains)) {
+    if (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(coeff[i - 2], walkfn).has_value()) {
       return ffi::Array<PrimExpr>();
     }
   }
diff --git a/src/arith/int_set.cc b/src/arith/int_set.cc
index 7c63eb908f09..9063c0b6208c 100644
--- a/src/arith/int_set.cc
+++ b/src/arith/int_set.cc
@@ -24,6 +24,7 @@
 #include <tvm/arith/int_set.h>
 #include <tvm/arith/iter_affine_map.h>
 #include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
 #include <tvm/ffi/function.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/ir/prim/expr.h>
@@ -599,10 +600,13 @@ class IntervalSetEvaluator : public ExprFunctor<IntervalSet(const Expr&)> {
     }
     // If the indices do not contain any variables to be relaxed, return the TensorLoad itself.
     // Otherwise return `IntervalSet::everything()` since we have no knowledge on the buffer data.
+    auto walkfn = [dom_map = &this->dom_map_](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+      return dom_map->find(var) != dom_map->end()
+                 ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                 : ffi::WalkResult::Advance();
+    };
     for (const PrimExpr& index : op->indices) {
-      if (UsesVar(index, [dom_map = &this->dom_map_](const VarNode* var) {
-            return dom_map->find(ffi::GetRef<Var>(var)) != dom_map->end();
-          })) {
+      if (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(index, walkfn).has_value()) {
         return IntervalSet::Everything();
       }
     }
diff --git a/src/arith/ir_mutator_with_analyzer.h b/src/arith/ir_mutator_with_analyzer.h
index da331ada86aa..d05e8e9a5f86 100644
--- a/src/arith/ir_mutator_with_analyzer.h
+++ b/src/arith/ir_mutator_with_analyzer.h
@@ -26,6 +26,7 @@
 
 #include <tvm/arith/analyzer.h>
 #include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
 #include <tvm/ir/scope_stack.h>
 #include <tvm/ir/with_context.h>
 #include <tvm/tirx/analysis.h>
@@ -108,8 +109,12 @@ class IRMutatorWithAnalyzer : public tirx::StmtExprMutator {
     auto f_use_itervar = [&iter_var_nodes](const tirx::VarNode* v) {
       return iter_var_nodes.count(v);
     };
+    auto walkfn = [&](const tirx::Var& var) -> ffi::Expected<ffi::WalkResult> {
+      return f_use_itervar(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                      : ffi::WalkResult::Advance();
+    };
     // simple heuristics for detecting predicate
-    if (tirx::UsesVar(condition, f_use_itervar)) {
+    if (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(condition, walkfn).has_value()) {
       iter_predicates_.push_back(condition);
       callback();
       iter_predicates_.pop_back();
diff --git a/src/arith/iter_affine_map.cc b/src/arith/iter_affine_map.cc
index 3b8daaa7f4d8..534b826d61f0 100644
--- a/src/arith/iter_affine_map.cc
+++ b/src/arith/iter_affine_map.cc
@@ -23,6 +23,7 @@
 #include <tvm/arith/analyzer.h>
 #include <tvm/arith/iter_affine_map.h>
 #include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/ir/prim/expr.h>
 #include <tvm/tirx/analysis.h>
@@ -1348,12 +1349,20 @@ bool MatchBoundConstraints(PrimExpr pred, ffi::Map<PrimVar, Range>* input_iters,
     auto f_use_itervar = [&input_iter_nodes](const VarNode* v) {
       return input_iter_nodes.count(v);
     };
+    auto walkfn = [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+      return f_use_itervar(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                      : ffi::WalkResult::Advance();
+    };
+    bool lhs_uses_itervar =
+        ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(lhs_expr, walkfn).has_value();
+    bool rhs_uses_itervar =
+        ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(rhs_expr, walkfn).has_value();
     bool bound_at_left;
-    if (UsesVar(lhs_expr, f_use_itervar) || UsesVar(rhs_expr, f_use_itervar)) {
+    if (lhs_uses_itervar || rhs_uses_itervar) {
       // At least it uses one input iter
-      if (is_const_int(lhs_expr) || !UsesVar(lhs_expr, f_use_itervar)) {
+      if (is_const_int(lhs_expr) || !lhs_uses_itervar) {
         bound_at_left = true;
-      } else if (is_const_int(rhs_expr) || !UsesVar(rhs_expr, f_use_itervar)) {
+      } else if (is_const_int(rhs_expr) || !rhs_uses_itervar) {
         bound_at_left = false;
       } else {
         bound_at_left = false;  // accumulate bound to rhs
@@ -1361,14 +1370,14 @@ bool MatchBoundConstraints(PrimExpr pred, ffi::Map<PrimVar, Range>* input_iters,
         lhs_expr = 0;
         rhs_expr = 0;
         std::function<void(const PrimExpr&, bool)> f_extract =
-            [&lhs_expr, &rhs_expr, f_use_itervar, &f_extract](const PrimExpr& part, bool sign) {
+            [&lhs_expr, &rhs_expr, &walkfn, &f_extract](const PrimExpr& part, bool sign) {
               if (const prim::AddNode* add = part.as<prim::AddNode>()) {
                 f_extract(add->a, sign);
                 f_extract(add->b, sign);
               } else if (const prim::SubNode* sub = part.as<prim::SubNode>()) {
                 f_extract(sub->a, sign);
                 f_extract(sub->b, !sign);
-              } else if (UsesVar(part, f_use_itervar)) {
+              } else if (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(part, walkfn).has_value()) {
                 lhs_expr = sign ? lhs_expr + part : lhs_expr - part;
               } else {
                 rhs_expr = sign ? rhs_expr - part : rhs_expr + part;
@@ -1429,8 +1438,15 @@ bool IterRangeSanityCheck(const ffi::Map<PrimVar, Range>& iter_ranges) {
   std::unordered_set<Var> iters;
   for (const auto& it : iter_ranges) iters.insert(it.first);
   auto f = [&](const VarNode* var) { return iters.count(ffi::GetRef<Var>(var)); };
+  auto walkfn = [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+    return f(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                        : ffi::WalkResult::Advance();
+  };
   for (const auto& it : iter_ranges) {
-    if (UsesVar(it.second->min, f) || UsesVar(it.second->extent, f)) return false;
+    if (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(it.second->min, walkfn).has_value() ||
+        ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(it.second->extent, walkfn).has_value()) {
+      return false;
+    }
   }
   return true;
 }
diff --git a/src/relax/analysis/tir_op_pattern_kind.cc b/src/relax/analysis/tir_op_pattern_kind.cc
index 5287e56d96dc..341d7534cfc2 100644
--- a/src/relax/analysis/tir_op_pattern_kind.cc
+++ b/src/relax/analysis/tir_op_pattern_kind.cc
@@ -19,6 +19,7 @@
 
 #include <tvm/arith/iter_affine_map.h>
 #include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/relax/analysis.h>
 #include <tvm/relax/op_attr_types.h>
@@ -235,10 +236,13 @@ class PatternKindAnalyzer : public StmtExprVisitor {
         return false;
       }
     }
+    auto walkfn = [&vars](const tirx::Var& var) -> ffi::Expected<ffi::WalkResult> {
+      return !vars.count(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                    : ffi::WalkResult::Advance();
+    };
     for (const PrimExpr& load_index : load->indices) {
       // return false if there are vars used in load indices but not in store indices.
-      if (tirx::UsesVar(load_index,
-                        [&vars](const tirx::VarNode* var) { return !vars.count(var); })) {
+      if (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(load_index, walkfn).has_value()) {
         return false;
       }
     }
@@ -318,17 +322,20 @@ class PatternKindAnalyzer : public StmtExprVisitor {
    */
   static bool IsPureReducePattern(ffi::Array<tirx::Var> reduce_loops,
                                   ffi::Array<PrimExpr> indices) {
+    auto walkfn = [&](const tirx::Var& var) -> ffi::Expected<ffi::WalkResult> {
+      return std::any_of(reduce_loops.begin(), reduce_loops.end(),
+                         [&](const tirx::Var& loop) { return loop.same_as(var); })
+                 ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                 : ffi::WalkResult::Advance();
+    };
     for (const PrimExpr& e : indices) {
-      int id = -1;
-      if (UsesVar(e, [&](const tirx::VarNode* var) {
-            for (size_t i = 0; i < reduce_loops.size(); ++i) {
-              if (reduce_loops[i].get() == var) {
-                id = i;
-                return true;
-              }
-            }
-            return false;
-          })) {
+      auto result = ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(e, walkfn);
+      if (result.has_value()) {
+        tirx::Var var = result.value()->value.cast<tirx::Var>();
+        int id =
+            std::distance(reduce_loops.begin(),
+                          std::find_if(reduce_loops.begin(), reduce_loops.end(),
+                                       [&](const tirx::Var& loop) { return loop.same_as(var); }));
         if (!reduce_loops[id].same_as(e)) {
           return false;
         }
diff --git a/src/relax/transform/rewrite_dataflow_reshape.cc b/src/relax/transform/rewrite_dataflow_reshape.cc
index 464be7d774cf..e99f10219d2a 100644
--- a/src/relax/transform/rewrite_dataflow_reshape.cc
+++ b/src/relax/transform/rewrite_dataflow_reshape.cc
@@ -22,6 +22,7 @@
  */
 #include <tvm/arith/analyzer.h>
 #include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/relax/analysis.h>
 #include <tvm/relax/expr_functor.h>
@@ -41,8 +42,11 @@ std::vector<size_t> GetUsedTensorArgIndices(const tirx::PrimFunc& fn, size_t num
   for (size_t i = 0; i < num_args; ++i) {
     if (auto buffer = fn->params[i].as<tirx::BufferVar>()) {
       auto buffer_var = buffer.value().var();
-      if (tirx::UsesVar(fn->body,
-                        [=](const tirx::VarNode* var) { return var == buffer_var.get(); })) {
+      auto walkfn = [=](const tirx::Var& var) -> ffi::Expected<ffi::WalkResult> {
+        return var.get() == buffer_var.get() ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                             : ffi::WalkResult::Advance();
+      };
+      if (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(fn->body, walkfn).has_value()) {
         indices.push_back(i);
       }
     }
diff --git a/src/s_tir/meta_schedule/postproc/rewrite_reduction_block.cc b/src/s_tir/meta_schedule/postproc/rewrite_reduction_block.cc
index 92ab70201d70..9a7a4acf6e7a 100644
--- a/src/s_tir/meta_schedule/postproc/rewrite_reduction_block.cc
+++ b/src/s_tir/meta_schedule/postproc/rewrite_reduction_block.cc
@@ -16,6 +16,7 @@
  * specific language governing permissions and limitations
  * under the License.
  */
+#include <tvm/ffi/extra/structural_visit.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/s_tir/stmt.h>
 
@@ -67,6 +68,10 @@ struct ReductionBlockFinder : private StmtVisitor {
       return true;
     }
     auto f_find = [this](const VarNode* var) -> bool { return thread_bound_loop_vars_.count(var); };
+    auto walkfn = [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+      return f_find(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                               : ffi::WalkResult::Advance();
+    };
     const SBlockNode* block = realize->block.get();
     TVM_FFI_ICHECK_EQ(block->iter_vars.size(), realize->iter_values.size());
     int n = block->iter_vars.size();
@@ -74,7 +79,7 @@ struct ReductionBlockFinder : private StmtVisitor {
       IterVar iter_var = block->iter_vars[i];
       PrimExpr binding = realize->iter_values[i];
       if (iter_var->iter_type == tirx::kCommReduce) {
-        if (UsesVar(binding, f_find)) {
+        if (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(binding, walkfn).has_value()) {
           return false;
         }
       }
diff --git a/src/s_tir/schedule/analysis/analysis.cc b/src/s_tir/schedule/analysis/analysis.cc
index 229c03e1ca04..20776081f4c5 100644
--- a/src/s_tir/schedule/analysis/analysis.cc
+++ b/src/s_tir/schedule/analysis/analysis.cc
@@ -17,6 +17,7 @@
  * under the License.
  */
 #include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/ir/op.h>
 #include <tvm/s_tir/stmt.h>
@@ -1826,6 +1827,14 @@ ffi::Optional<TensorizeInfo> GetTensorizeLoopMapping(const s_tir::ScheduleState&
   //       C[i, j] += A[i, k] * B[k, j]
 
   int next_block_ind = block_loops.size() - 1;
+  auto desc_walkfn = [&desc_loop_vars](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+    return desc_loop_vars.count(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                           : ffi::WalkResult::Advance();
+  };
+  auto block_walkfn = [&block_loop_vars](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+    return block_loop_vars.count(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                            : ffi::WalkResult::Advance();
+  };
   for (int i_desc = n_desc_vars - 1; i_desc >= 0; --i_desc) {
     // Step 3.1. Find the corresponding loop of the i_desc-th block var of desc
     const PrimExpr& desc_bind = desc_block->iter_values[i_desc];
@@ -1834,8 +1843,7 @@ ffi::Optional<TensorizeInfo> GetTensorizeLoopMapping(const s_tir::ScheduleState&
     for (int i = 0, n = desc_loops.size(); i < n; ++i) {
       // Check if desc_bind = loops[i]->loop_var + stuff-irrelevant-of-loop-vars
       PrimExpr residual = analyzer->Simplify(desc_bind - desc_loops[i]->loop_var);
-      if (!UsesVar(residual,
-                   [&desc_loop_vars](const VarNode* var) { return desc_loop_vars.count(var); })) {
+      if (!ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(residual, desc_walkfn).has_value()) {
         desc_loop = desc_loops[i];
         iter_type_desc = iter_types_desc[i];
         break;
@@ -1869,8 +1877,7 @@ ffi::Optional<TensorizeInfo> GetTensorizeLoopMapping(const s_tir::ScheduleState&
       if (ret->loop_map.find(block_loop_sref) != ret->loop_map.end()) continue;
 
       PrimExpr residual = analyzer->Simplify(block_bind - block_loops[i]->loop_var);
-      if (UsesVar(residual,
-                  [&block_loop_vars](const VarNode* var) { return block_loop_vars.count(var); })) {
+      if (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(residual, block_walkfn).has_value()) {
         continue;
       }
       // padding is allowed only when the block has trivial bindings
diff --git a/src/s_tir/schedule/analysis/reducer.cc b/src/s_tir/schedule/analysis/reducer.cc
index ac15128f7c24..c122ea66a091 100644
--- a/src/s_tir/schedule/analysis/reducer.cc
+++ b/src/s_tir/schedule/analysis/reducer.cc
@@ -17,6 +17,7 @@
  * under the License.
  */
 #include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
 #include <tvm/te/operation.h>
 
 #include "../utils.h"
@@ -553,10 +554,13 @@ bool ReductionIterNotIndexOutputBuffer(const SBlock& block) {
     buffer_allocated.insert(buffer.get());
   }
 
+  auto walkfn = [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+    return reduction_block_iters.count(var.get())
+               ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+               : ffi::WalkResult::Advance();
+  };
   auto f_uses_reduction_block_var = [&](const PrimExpr& expr) -> bool {
-    return UsesVar(expr, [&](const VarNode* var) {  //
-      return reduction_block_iters.count(var);
-    });
+    return ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(expr, walkfn).has_value();
   };
 
   std::unordered_map<const VarNode*, const VarNode*> match_buffer_sources;
diff --git a/src/s_tir/schedule/primitive/blockize_tensorize.cc b/src/s_tir/schedule/primitive/blockize_tensorize.cc
index ca5053c92134..3dea0688496f 100644
--- a/src/s_tir/schedule/primitive/blockize_tensorize.cc
+++ b/src/s_tir/schedule/primitive/blockize_tensorize.cc
@@ -18,6 +18,7 @@
  */
 
 #include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
 #include <tvm/runtime/logging.h>
 
 #include <functional>
@@ -32,11 +33,6 @@ namespace s_tir {
 using namespace tvm::prim;
 using namespace tvm::tirx;
 
-template <class T>
-bool UsesVar(const T& x, const Var& var) {
-  return tirx::UsesVar(x, [tgt = var.get()](const VarNode* v) { return v == tgt; });
-}
-
 Range RangeFromExtent(const PrimExpr& extent) {
   return Range::FromMinExtent(IntImm(extent.ty(), 0), extent);
 }
@@ -110,9 +106,11 @@ ffi::Array<ffi::Array<arith::IterMark>> TrivialSubspaceDivision(
       var_set.insert(var.get());
     }
     return [var_set = std::move(var_set)](const PrimExpr& expr) -> bool {
-      return tirx::UsesVar(expr, [&var_set](const VarNode* var) {
-        return var_set.count(var);  //
-      });
+      auto walkfn = [&var_set](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+        return var_set.count(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                        : ffi::WalkResult::Advance();
+      };
+      return ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(expr, walkfn).has_value();
     };
   };
   auto use_outer_loop_vars = make_uses_var(outer_iters);
@@ -335,8 +333,13 @@ Stmt GenerateOuterInit(const Stmt& block_init, const SBlockRealize& inner_realiz
   for (int i = 0; i < n; ++i) {
     const IterVar& old_iter_var = inner_block->iter_vars[i];
     const PrimExpr& iter_value = inner_realize->iter_values[i];
+    auto walkfn = [target =
+                       old_iter_var->var.get()](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+      return var.get() == target ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                 : ffi::WalkResult::Advance();
+    };
     if (old_iter_var->iter_type == IterVarType::kDataPar &&
-        UsesVar(block_init, old_iter_var->var)) {
+        ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(block_init, walkfn).has_value()) {
       ffi::ObjectPtr<IterVarNode> new_iter_var = ffi::make_object<IterVarNode>(*old_iter_var.get());
       new_iter_var->var = new_iter_var->var.CopyWithSuffix("_init");
       subst_map.Set(old_iter_var->var, new_iter_var->var);
@@ -358,8 +361,13 @@ Stmt GenerateOuterInit(const Stmt& block_init, const SBlockRealize& inner_realiz
   // Step 3. Create the loop nest on top of the block
   for (const ForNode* loop : loops) {
     bool is_init_loop = false;
+    auto walkfn = [target =
+                       loop->loop_var.get()](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+      return var.get() == target ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                 : ffi::WalkResult::Advance();
+    };
     for (const PrimExpr& init_binding : iter_values) {
-      if (UsesVar(init_binding, loop->loop_var)) {
+      if (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(init_binding, walkfn).has_value()) {
         is_init_loop = true;
         break;
       }
diff --git a/src/s_tir/schedule/primitive/for_kind.cc b/src/s_tir/schedule/primitive/for_kind.cc
index c2c7c1c887f3..b23d50d0f41f 100644
--- a/src/s_tir/schedule/primitive/for_kind.cc
+++ b/src/s_tir/schedule/primitive/for_kind.cc
@@ -17,6 +17,7 @@
  * under the License.
  */
 #include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
 
 #include "../utils.h"
 
@@ -97,7 +98,11 @@ void CheckLoopParallelizableInBlock(const ScheduleState& self, ForKind for_kind,
     const IterVar& iter_var = block->iter_vars[i];
     const PrimExpr& binding = block_realize->iter_values[i];
 
-    if (!UsesVar(binding, [v = loop_var.get()](const VarNode* var) { return var == v; })) {
+    auto walkfn = [v = loop_var.get()](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+      return var.get() == v ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                            : ffi::WalkResult::Advance();
+    };
+    if (!ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(binding, walkfn).has_value()) {
       continue;
     }
     // Only two cases are allowed:
diff --git a/src/s_tir/schedule/primitive/loop_transformation.cc b/src/s_tir/schedule/primitive/loop_transformation.cc
index e7ee773ac58d..6c354769df34 100644
--- a/src/s_tir/schedule/primitive/loop_transformation.cc
+++ b/src/s_tir/schedule/primitive/loop_transformation.cc
@@ -17,6 +17,7 @@
  * under the License.
  */
 #include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
 
 #include "../utils.h"
 
@@ -907,15 +908,13 @@ StmtSRef Fuse(ScheduleState self, const ffi::Array<StmtSRef>& loop_srefs,
     outer_loop_sref = sref;
     outer_loop = loop;
     CheckLoopStartsWithZero(self, sref, analyzer.get());
-    const VarNode* used_var = nullptr;
-    auto f_contain = [&outer_loop_vars, &used_var](const VarNode* var) {
-      if (outer_loop_vars.count(var)) {
-        used_var = var;
-        return true;
-      }
-      return false;
+    auto walkfn = [&outer_loop_vars](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+      return outer_loop_vars.count(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                              : ffi::WalkResult::Advance();
     };
-    if (UsesVar(loop->extent, f_contain)) {
+    auto result = ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(loop->extent, walkfn);
+    if (result.has_value()) {
+      Var used_var = result.value()->value.cast<Var>();
       throw DependentLoopError(self->mod, ffi::GetRef<For>(loop), used_var->name,
                                DependentLoopError::PrimitiveKind::kFuse);
     }
@@ -1105,15 +1104,16 @@ For ConstructNewLoopChain(const ScheduleState& self, std::vector<const StmtSRefN
     } else {
       n->body = loop_sref->StmtAs<ForNode>()->body;
     }
-    const VarNode* used_var = nullptr;
-    auto f_contain = [&inner_vars, &used_var](const VarNode* var) {
-      if (inner_vars.count(var)) {
-        used_var = var;
-        return true;
-      }
-      return false;
+    auto walkfn = [&inner_vars](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+      return inner_vars.count(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                         : ffi::WalkResult::Advance();
     };
-    if (UsesVar(copy->min, f_contain) || UsesVar(copy->extent, f_contain)) {
+    auto result = ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(copy->min, walkfn);
+    if (!result.has_value()) {
+      result = ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(copy->extent, walkfn);
+    }
+    if (result.has_value()) {
+      Var used_var = result.value()->value.cast<Var>();
       throw DependentLoopError(self->mod, ffi::GetRef<For>(copy), used_var->name,
                                DependentLoopError::PrimitiveKind::kReorder);
     }
diff --git a/src/s_tir/schedule/primitive/reduction.cc b/src/s_tir/schedule/primitive/reduction.cc
index beb05afbd542..a96face461b2 100644
--- a/src/s_tir/schedule/primitive/reduction.cc
+++ b/src/s_tir/schedule/primitive/reduction.cc
@@ -17,6 +17,7 @@
  * under the License.
  */
 #include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/te/operation.h>
 
@@ -128,7 +129,11 @@ class LoopHeightError : public ScheduleError {
         }
         // loop_var of a higher loop shouldn't contain loop var
         const Var& loop_var = higher_loop->StmtAs<ForNode>()->loop_var;
-        if (UsesVar(binding, [v = loop_var.get()](const VarNode* var) { return var == v; })) {
+        auto walkfn = [v = loop_var.get()](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+          return var.get() == v ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                : ffi::WalkResult::Advance();
+        };
+        if (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(binding, walkfn).has_value()) {
           const ForNode* loop = TVM_SREF_TO_FOR(loop_sref);
           throw LoopHeightError(mod, ffi::GetRef<For>(loop), ffi::GetRef<SBlock>(block));
         }
@@ -169,7 +174,13 @@ PrimExpr RewriteInitPredicate(PrimExpr pred,
   auto uses_discarded_loop = [&discarded_loops](const VarNode* var) {
     return discarded_loops.count(var);
   };
-  return UsesVar(pred, uses_discarded_loop) ? IntImm::Bool(true) : pred;
+  auto walkfn = [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+    return uses_discarded_loop(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                          : ffi::WalkResult::Advance();
+  };
+  return ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(pred, walkfn).has_value()
+             ? IntImm::Bool(true)
+             : pred;
 }
 
 StmtSRef DecomposeReduction(ScheduleState self, const StmtSRef& block_sref,
@@ -244,8 +255,12 @@ StmtSRef DecomposeReduction(ScheduleState self, const StmtSRef& block_sref,
   for (int i = static_cast<int>(loops.size()) - 1; i >= 0; --i) {
     const VarNode* loop_var = loops[i]->StmtAs<ForNode>()->loop_var.get();
     bool discarded = true;
+    auto walkfn = [v = loop_var](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+      return var.get() == v ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                            : ffi::WalkResult::Advance();
+    };
     for (const PrimExpr& expr : init_realize->iter_values) {
-      if (!UsesVar(expr, [v = loop_var](const VarNode* var) { return var == v; })) {
+      if (!ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(expr, walkfn).has_value()) {
         continue;
       }
       // The loop is related to init block bindings;
@@ -896,9 +911,12 @@ class RFactorBlockCreator : public BaseBlockCreator {
   void CreateNormalIters(int idx) final {
     IterVar old_iter = old_block_realize_->block->iter_vars[idx];
     PrimExpr old_binding = old_block_realize_->iter_values[idx];
+    auto walkfn = [v = rf_loop_->loop_var.get()](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+      return var.get() == v ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                            : ffi::WalkResult::Advance();
+    };
     if (old_iter->iter_type == IterVarType::kDataPar ||
-        !UsesVar(old_binding,
-                 [v = rf_loop_->loop_var.get()](const VarNode* var) { return var == v; })) {
+        !ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(old_binding, walkfn).has_value()) {
       // The old block iter is either a data parallel block iter, or a reduction block iter that
       // doesn't touch the rfactor loop. In this case reuse the old reduction block iter and its
       // corresponding binding.
diff --git a/src/s_tir/transform/compact_buffer_region.cc b/src/s_tir/transform/compact_buffer_region.cc
index 2ccbd8e905be..7255317fa358 100644
--- a/src/s_tir/transform/compact_buffer_region.cc
+++ b/src/s_tir/transform/compact_buffer_region.cc
@@ -25,6 +25,7 @@
 #include <tvm/arith/int_set.h>
 #include <tvm/arith/int_solver.h>
 #include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/s_tir/stmt.h>
 #include <tvm/s_tir/transform.h>
@@ -472,7 +473,11 @@ class BufferAccessRegionCollector : public StmtExprVisitor {
         return std::any_of(ancestor_iters_.begin(), ancestor_iters_.end(),
                            [v](const IterVar& n) { return n->var.get() == v; });
       };
-      if (UsesVar(extent, is_loop_var)) {
+      auto walkfn = [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+        return is_loop_var(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                      : ffi::WalkResult::Advance();
+      };
+      if (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(extent, walkfn).has_value()) {
         // try estimate a constant upperbound on region's extent
         int64_t upperbound = dom_analyzer_->const_int_bound(extent)->max_value;
         if (upperbound != arith::ConstIntBound::kPosInf) {
diff --git a/src/s_tir/transform/hoist_expression.cc b/src/s_tir/transform/hoist_expression.cc
index 6aa8082a98b0..40a4b47daf3c 100644
--- a/src/s_tir/transform/hoist_expression.cc
+++ b/src/s_tir/transform/hoist_expression.cc
@@ -22,6 +22,7 @@
  */
 #include <tvm/arith/analyzer.h>
 #include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
 #include <tvm/ffi/function.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/ir/prim/expr.h>
@@ -217,9 +218,14 @@ class HoistInfoCollector : public StmtExprVisitor {
     if (auto info = FindHoistDestination(cond)) {
       if (!info->reached_sequential_node) {
         // Record whether this conditional uses any block variables.
-        bool uses_block_var = active_block_vars.size() && UsesVar(cond, [&](const VarNode* var) {
-                                return active_block_vars.count(var);
-                              });
+        auto walkfn = [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+          return active_block_vars.count(var.get())
+                     ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                     : ffi::WalkResult::Advance();
+        };
+        bool uses_block_var =
+            active_block_vars.size() &&
+            ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(cond, walkfn).has_value();
 
         std::unordered_set<const VarNode*> let_bindings_used;
 
@@ -389,18 +395,16 @@ class HoistInfoCollector : public StmtExprVisitor {
 
     for (auto it = active_loops.rbegin(); it != active_loops.rend(); it++) {
       Var loop_var = it->loop_var;
-      bool uses_loop_var = UsesVar(expr, [&](const VarNode* var) -> bool {
-        if (var == loop_var.get()) {
-          return true;
+      auto walkfn = [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+        bool matches = var.get() == loop_var.get();
+        if (!matches) {
+          auto it = let_var_to_loop_vars.find(var.get());
+          matches = it != let_var_to_loop_vars.end() && it->second.count(loop_var.get());
         }
-
-        auto it = let_var_to_loop_vars.find(var);
-        if (it == let_var_to_loop_vars.end()) {
-          return false;
-        }
-
-        return it->second.count(loop_var.get());
-      });
+        return matches ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                       : ffi::WalkResult::Advance();
+      };
+      bool uses_loop_var = ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(expr, walkfn).has_value();
 
       bool is_disabled_hoist_across_block_var =
           !config->FlagSet(HoistedConditionals::kUsingBlockVar) && it->IsBlockVariable();
diff --git a/src/s_tir/transform/loop_partition.cc b/src/s_tir/transform/loop_partition.cc
index 56d6b1303895..e04a48687e40 100644
--- a/src/s_tir/transform/loop_partition.cc
+++ b/src/s_tir/transform/loop_partition.cc
@@ -23,6 +23,7 @@
 #include <tvm/arith/analyzer.h>
 #include <tvm/arith/bound.h>
 #include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
 #include <tvm/ffi/function.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/ir/prim/builtin.h>
@@ -248,7 +249,14 @@ class PartitionFinder : public StmtExprVisitor {
 
   void VisitStmt_(const ForNode* op) final {
     auto f_vset_contains = [this](const VarNode* var) { return out_vars_.count(var); };
-    if (UsesVar(op->min, f_vset_contains) || UsesVar(op->extent, f_vset_contains)) return;
+    auto walkfn = [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+      return f_vset_contains(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                        : ffi::WalkResult::Advance();
+    };
+    if (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(op->min, walkfn).has_value() ||
+        ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(op->extent, walkfn).has_value()) {
+      return;
+    }
 
     const VarNode* var = op->loop_var.get();
     hint_map_.insert({var, IntSet::Interval(op->min, op->min + op->extent - 1)});
@@ -299,7 +307,11 @@ class PartitionFinder : public StmtExprVisitor {
     // For cond, find out the interval, if exists, in which we can prove that cond is
     // true. Also find the interval, if exists, in which we can prove that cond is
     // false.
-    if (UsesVar(cond, [this](const VarNode* var) { return var == current_var_.get(); })) {
+    auto walkfn = [this](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+      return var.get() == current_var_.get() ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                             : ffi::WalkResult::Advance();
+    };
+    if (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(cond, walkfn).has_value()) {
       IntSet interval =
           DeduceBound(current_var_.as_or_throw<PrimExpr>(), cond, hint_map_, relax_map_);
       if (!interval.IsNothing()) {
diff --git a/src/s_tir/transform/thread_storage_sync.cc b/src/s_tir/transform/thread_storage_sync.cc
index d284accd42ef..200ef98ae932 100644
--- a/src/s_tir/transform/thread_storage_sync.cc
+++ b/src/s_tir/transform/thread_storage_sync.cc
@@ -20,6 +20,7 @@
 /*!
  * \file thread_storage_sync.cc
  */
+#include <tvm/ffi/extra/structural_visit.h>
 #include <tvm/ffi/function.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/ir/prim/builtin.h>
@@ -237,9 +238,15 @@ class ThreadSyncPlanner : public StorageAccessVisitor {
           auto f_uses_thread_index = [=](const tvm::tirx::VarNode* parameter) {
             return parameter == thread_index_var;
           };
-          depends_on_thread_index = depends_on_thread_index &&
-                                    UsesVar(curr_index, f_uses_thread_index) &&
-                                    UsesVar(prev_index, f_uses_thread_index);
+          auto walkfn = [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+            return f_uses_thread_index(var.get())
+                       ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                       : ffi::WalkResult::Advance();
+          };
+          depends_on_thread_index =
+              depends_on_thread_index &&
+              ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(curr_index, walkfn).has_value() &&
+              ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(prev_index, walkfn).has_value();
         }
       } else {
         has_same_index = false;
diff --git a/src/tirx/analysis/var_touch.cc b/src/tirx/analysis/var_touch.cc
deleted file mode 100644
index d3441b3351bc..000000000000
--- a/src/tirx/analysis/var_touch.cc
+++ /dev/null
@@ -1,80 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements.  See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership.  The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License.  You may obtain a copy of the License at
- *
- *   http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied.  See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-/*!
- * \file var_touch.cc
- * \brief Implementation of simple passes
- */
-#include <tvm/tirx/analysis.h>
-#include <tvm/tirx/stmt_functor.h>
-
-namespace tvm {
-namespace tirx {
-
-class VarTouchVisitor : public StmtExprVisitor {
- public:
-  explicit VarTouchVisitor(std::function<bool(const VarNode*)> var_set)
-      : var_set_(std::move(var_set)) {}
-
-  void VisitStmt(const Stmt& stmt) final {
-    if (use_var_) return;
-    StmtExprVisitor::VisitStmt(stmt);
-  }
-
-  void VisitExpr(const Expr& e) final {
-    if (use_var_) return;
-    StmtExprVisitor::VisitExpr(e);
-  }
-
-  void VisitExpr_(const VarNode* op) final { Handle(op); }
-
-  void VisitStmt_(const BufferStoreNode* op) final {
-    Handle(op->buffer.get());
-    StmtVisitor::VisitStmt_(op);
-  }
-
-  void VisitExpr_(const TensorLoadNode* op) final {
-    Handle(op->source.as_or_throw<tvm::tirx::BufferVar>().get());
-    ExprVisitor::VisitExpr_(op);
-  }
-
-  void Handle(const VarNode* var) {
-    if (var_set_(var)) use_var_ = true;
-  }
-
-  bool use_var_{false};
-
- private:
-  std::function<bool(const VarNode*)> var_set_;
-};
-
-bool UsesVar(const Stmt& stmt, std::function<bool(const VarNode*)> var_set) {
-  VarTouchVisitor visitor(std::move(var_set));
-  visitor(stmt);
-  return visitor.use_var_;
-}
-
-bool UsesVar(const PrimExpr& expr, std::function<bool(const VarNode*)> var_set) {
-  VarTouchVisitor visitor(std::move(var_set));
-  visitor(expr);
-  return visitor.use_var_;
-}
-
-}  // namespace tirx
-}  // namespace tvm
diff --git a/src/tirx/script/printer/for_loop.cc b/src/tirx/script/printer/for_loop.cc
index 363a2d466f0b..4d0f2016cdec 100644
--- a/src/tirx/script/printer/for_loop.cc
+++ b/src/tirx/script/printer/for_loop.cc
@@ -16,6 +16,8 @@
  * specific language governing permissions and limitations
  * under the License.
  */
+#include <tvm/ffi/extra/structural_visit.h>
+
 #include "./utils.h"
 
 namespace tvm {
@@ -28,9 +30,12 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
       std::vector<const tirx::ForNode*> grid;
       std::unordered_set<const tirx::VarNode*> grid_loop_vars;
       auto f_var_dep = [&grid_loop_vars](const PrimExpr& e) -> bool {
-        return tirx::UsesVar(e, [&grid_loop_vars](const tirx::VarNode* v) -> bool {  //
-          return grid_loop_vars.count(v);
-        });
+        auto walkfn = [&grid_loop_vars](const tirx::Var& var) -> ffi::Expected<ffi::WalkResult> {
+          return grid_loop_vars.count(var.get())
+                     ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                     : ffi::WalkResult::Advance();
+        };
+        return ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(e, walkfn).has_value();
       };
       if (d->cfg->syntax_sugar) {
         for (const tirx::ForNode* l = loop.get(); l != nullptr; l = l->body.as<tirx::ForNode>()) {
diff --git a/src/tirx/transform/ir_utils.cc b/src/tirx/transform/ir_utils.cc
index e41223d819fc..5633ba59865b 100644
--- a/src/tirx/transform/ir_utils.cc
+++ b/src/tirx/transform/ir_utils.cc
@@ -26,6 +26,7 @@
 #include <tvm/arith/analyzer.h>
 #include <tvm/arith/int_solver.h>
 #include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/ir/scope_stack.h>
 #include <tvm/s_tir/stmt.h>
@@ -566,7 +567,12 @@ class IRConvertSSA final : public StmtExprMutator {
     if (buffer.get() == var) return true;
 
     auto uses_var = [var](const PrimExpr& expr) {
-      return expr.defined() && UsesVar(expr, [var](const VarNode* node) { return node == var; });
+      auto walkfn = [var](const Var& candidate) -> ffi::Expected<ffi::WalkResult> {
+        return candidate.get() == var ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(candidate))
+                                      : ffi::WalkResult::Advance();
+      };
+      return expr.defined() &&
+             ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(expr, walkfn).has_value();
     };
     if (uses_var(buffer->elem_offset)) return true;
     for (const PrimExpr& dim : buffer->shape) {
diff --git a/src/tirx/transform/lower_warp_memory.cc b/src/tirx/transform/lower_warp_memory.cc
index ed2afbfc2a86..4f268b40e096 100644
--- a/src/tirx/transform/lower_warp_memory.cc
+++ b/src/tirx/transform/lower_warp_memory.cc
@@ -28,6 +28,7 @@
 #include <tvm/arith/analyzer.h>
 #include <tvm/arith/pattern.h>
 #include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
 #include <tvm/ffi/function.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/ir/op.h>
@@ -384,8 +385,11 @@ class WarpAccessRewriter : protected StmtExprMutator {
 
     auto [local_index, group] = SplitIndexByGroup(op->indices[0]);
     // invariance: local index must do not contain warp id
-    TVM_FFI_ICHECK(
-        !UsesVar(local_index, [this](const VarNode* var) { return var == warp_index_.get(); }))
+    auto walkfn = [this](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+      return var.get() == warp_index_.get() ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                            : ffi::WalkResult::Advance();
+    };
+    TVM_FFI_ICHECK(!ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(local_index, walkfn).has_value())
         << "LowerWarpMemory failed to rewrite load to shuffle for index " << op->indices[0]
         << " local_index=" << local_index;