From f92b02415a47a787747d0f298e4bc742db050343 Mon Sep 17 00:00:00 2001
From: tqchen <tianqic@nvidia.com>
Date: Thu, 10 Sep 2026 02:36:53 +0000
Subject: [PATCH 1/4] [REFACTOR][TIR] Use StructuralWalk in UsesVar

---
 src/tirx/analysis/var_touch.cc        | 63 +++++++++------------------
 tests/cpp/tir_analysis_side_effect.cc | 26 +++++++++++
 2 files changed, 47 insertions(+), 42 deletions(-)

diff --git a/src/tirx/analysis/var_touch.cc b/src/tirx/analysis/var_touch.cc
index d3441b3351bc..32b760ffe843 100644
--- a/src/tirx/analysis/var_touch.cc
+++ b/src/tirx/analysis/var_touch.cc
@@ -21,59 +21,38 @@
  * \file var_touch.cc
  * \brief Implementation of simple passes
  */
+#include <tvm/ffi/extra/structural_visit.h>
 #include <tvm/tirx/analysis.h>
-#include <tvm/tirx/stmt_functor.h>
+
+#include <utility>
 
 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};
+namespace {
+
+template <typename T>
+bool UsesVarImpl(const T& value, std::function<bool(const VarNode*)> var_set) {
+  bool use_var = false;
+  ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+      value, [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+        if (var_set(var.get())) {
+          use_var = true;
+          return ffi::WalkResult::Interrupt();
+        }
+        return ffi::WalkResult::Advance();
+      });
+  return use_var;
+}
 
- private:
-  std::function<bool(const VarNode*)> var_set_;
-};
+}  // namespace
 
 bool UsesVar(const Stmt& stmt, std::function<bool(const VarNode*)> var_set) {
-  VarTouchVisitor visitor(std::move(var_set));
-  visitor(stmt);
-  return visitor.use_var_;
+  return UsesVarImpl(stmt, std::move(var_set));
 }
 
 bool UsesVar(const PrimExpr& expr, std::function<bool(const VarNode*)> var_set) {
-  VarTouchVisitor visitor(std::move(var_set));
-  visitor(expr);
-  return visitor.use_var_;
+  return UsesVarImpl(expr, std::move(var_set));
 }
 
 }  // namespace tirx
diff --git a/tests/cpp/tir_analysis_side_effect.cc b/tests/cpp/tir_analysis_side_effect.cc
index da55576861f2..4bfb53ad6a2b 100644
--- a/tests/cpp/tir_analysis_side_effect.cc
+++ b/tests/cpp/tir_analysis_side_effect.cc
@@ -35,3 +35,29 @@ TEST(SimplePasses, SideEffect) {
                                       .as_or_throw<PrimExpr>()) ==
                  tirx::CallEffectKind::kUpdateState);
 }
+
+TEST(SimplePasses, UsesVar) {
+  using namespace tvm;
+  using namespace tvm::tirx;
+
+  PrimVar i("i", PrimType::Int(32));
+  PrimVar j("j", PrimType::Int(32));
+  BufferVar buffer = decl_buffer({16}, PrimType::Float(32));
+  auto is_var = [](const Var& var) {
+    return [var](const VarNode* candidate) { return candidate == var.get(); };
+  };
+
+  EXPECT_TRUE(UsesVar(i + j, is_var(j)));
+  EXPECT_TRUE(UsesVar(BufferLoad(buffer, {i}), is_var(buffer)));
+  EXPECT_TRUE(UsesVar(BufferStore(buffer, FloatImm(PrimType::Float(32), 0), {j}), is_var(buffer)));
+
+  Stmt loop = For(i, 0, 4, ForKind::kSerial, Evaluate(j));
+  EXPECT_TRUE(UsesVar(loop, is_var(i)));
+
+  int visits = 0;
+  EXPECT_TRUE(UsesVar(i + j, [&](const VarNode* candidate) {
+    ++visits;
+    return candidate == i.get();
+  }));
+  EXPECT_EQ(visits, 1);
+}

From 7dbbf2013d00e8b3273a1a791c86ce33b30dd115 Mon Sep 17 00:00:00 2001
From: tqchen <tianqic@nvidia.com>
Date: Thu, 10 Sep 2026 12:28:44 +0000
Subject: [PATCH 2/4] [REFACTOR][TIR] Inline UsesVar structural walks

---
 src/tirx/analysis/var_touch.cc        | 32 +++++++++++----------------
 tests/cpp/tir_analysis_side_effect.cc | 26 ----------------------
 2 files changed, 13 insertions(+), 45 deletions(-)

diff --git a/src/tirx/analysis/var_touch.cc b/src/tirx/analysis/var_touch.cc
index 32b760ffe843..9d3475eeebf1 100644
--- a/src/tirx/analysis/var_touch.cc
+++ b/src/tirx/analysis/var_touch.cc
@@ -24,35 +24,29 @@
 #include <tvm/ffi/extra/structural_visit.h>
 #include <tvm/tirx/analysis.h>
 
-#include <utility>
-
 namespace tvm {
 namespace tirx {
 
-namespace {
-
-template <typename T>
-bool UsesVarImpl(const T& value, std::function<bool(const VarNode*)> var_set) {
-  bool use_var = false;
-  ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-      value, [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+bool UsesVar(const Stmt& stmt, std::function<bool(const VarNode*)> var_set) {
+  auto result = ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+      stmt, [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
         if (var_set(var.get())) {
-          use_var = true;
-          return ffi::WalkResult::Interrupt();
+          return ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var));
         }
         return ffi::WalkResult::Advance();
       });
-  return use_var;
-}
-
-}  // namespace
-
-bool UsesVar(const Stmt& stmt, std::function<bool(const VarNode*)> var_set) {
-  return UsesVarImpl(stmt, std::move(var_set));
+  return result.has_value();
 }
 
 bool UsesVar(const PrimExpr& expr, std::function<bool(const VarNode*)> var_set) {
-  return UsesVarImpl(expr, std::move(var_set));
+  auto result = ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+      expr, [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+        if (var_set(var.get())) {
+          return ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var));
+        }
+        return ffi::WalkResult::Advance();
+      });
+  return result.has_value();
 }
 
 }  // namespace tirx
diff --git a/tests/cpp/tir_analysis_side_effect.cc b/tests/cpp/tir_analysis_side_effect.cc
index 4bfb53ad6a2b..da55576861f2 100644
--- a/tests/cpp/tir_analysis_side_effect.cc
+++ b/tests/cpp/tir_analysis_side_effect.cc
@@ -35,29 +35,3 @@ TEST(SimplePasses, SideEffect) {
                                       .as_or_throw<PrimExpr>()) ==
                  tirx::CallEffectKind::kUpdateState);
 }
-
-TEST(SimplePasses, UsesVar) {
-  using namespace tvm;
-  using namespace tvm::tirx;
-
-  PrimVar i("i", PrimType::Int(32));
-  PrimVar j("j", PrimType::Int(32));
-  BufferVar buffer = decl_buffer({16}, PrimType::Float(32));
-  auto is_var = [](const Var& var) {
-    return [var](const VarNode* candidate) { return candidate == var.get(); };
-  };
-
-  EXPECT_TRUE(UsesVar(i + j, is_var(j)));
-  EXPECT_TRUE(UsesVar(BufferLoad(buffer, {i}), is_var(buffer)));
-  EXPECT_TRUE(UsesVar(BufferStore(buffer, FloatImm(PrimType::Float(32), 0), {j}), is_var(buffer)));
-
-  Stmt loop = For(i, 0, 4, ForKind::kSerial, Evaluate(j));
-  EXPECT_TRUE(UsesVar(loop, is_var(i)));
-
-  int visits = 0;
-  EXPECT_TRUE(UsesVar(i + j, [&](const VarNode* candidate) {
-    ++visits;
-    return candidate == i.get();
-  }));
-  EXPECT_EQ(visits, 1);
-}

From c2adec9e983fbfdb5bdb2a9804d1da0b937230a1 Mon Sep 17 00:00:00 2001
From: tqchen <tianqic@nvidia.com>
Date: Thu, 10 Sep 2026 13:25:22 +0000
Subject: [PATCH 3/4] [REFACTOR][TIR] Inline variable-use walks

---
 include/tvm/tirx/analysis.h                   | 16 ------
 src/arith/detect_linear_equation.cc           | 16 +++++-
 src/arith/int_set.cc                          | 12 +++--
 src/arith/ir_mutator_with_analyzer.h          |  9 +++-
 src/arith/iter_affine_map.cc                  | 49 +++++++++++++++--
 src/relax/analysis/tir_op_pattern_kind.cc     | 33 +++++++-----
 .../transform/rewrite_dataflow_reshape.cc     | 11 +++-
 .../postproc/rewrite_reduction_block.cc       |  9 +++-
 src/s_tir/schedule/analysis/analysis.cc       | 21 ++++++--
 src/s_tir/schedule/analysis/reducer.cc        | 12 +++--
 .../schedule/primitive/blockize_tensorize.cc  | 33 ++++++++----
 src/s_tir/schedule/primitive/for_kind.cc      |  9 +++-
 .../schedule/primitive/loop_transformation.cc | 35 ++++++------
 src/s_tir/schedule/primitive/reduction.cc     | 36 +++++++++++--
 src/s_tir/transform/compact_buffer_region.cc  |  9 +++-
 src/s_tir/transform/hoist_expression.cc       | 39 ++++++++------
 src/s_tir/transform/loop_partition.cc         | 29 +++++++++-
 src/s_tir/transform/thread_storage_sync.cc    | 22 ++++++--
 src/tirx/analysis/var_touch.cc                | 53 -------------------
 src/tirx/script/printer/for_loop.cc           | 13 +++--
 src/tirx/transform/ir_utils.cc                | 11 +++-
 src/tirx/transform/lower_warp_memory.cc       | 11 +++-
 22 files changed, 326 insertions(+), 162 deletions(-)
 delete mode 100644 src/tirx/analysis/var_touch.cc

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..0f97d2c96018 100644
--- a/src/arith/detect_linear_equation.cc
+++ b/src/arith/detect_linear_equation.cc
@@ -110,7 +110,13 @@ 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(); })) {
+    if (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+            e,
+            [this](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+              return var.get() == var_.get() ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                             : ffi::WalkResult::Advance();
+            })
+            .has_value()) {
       fail_ = true;
       return LinearEqEntry();
     } else {
@@ -161,7 +167,13 @@ ffi::Array<PrimExpr> DetectLinearEquation(const PrimExpr& e, const ffi::Array<Pr
   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],
+            [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+              return vset_contains(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                              : ffi::WalkResult::Advance();
+            })
+            .has_value()) {
       return ffi::Array<PrimExpr>();
     }
   }
diff --git a/src/arith/int_set.cc b/src/arith/int_set.cc
index 7c63eb908f09..3aba438f550e 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>
@@ -600,9 +601,14 @@ 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.
     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,
+              [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();
+              })
+              .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..8c40aa24fa4f 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>
@@ -109,7 +110,13 @@ class IRMutatorWithAnalyzer : public tirx::StmtExprMutator {
       return iter_var_nodes.count(v);
     };
     // simple heuristics for detecting predicate
-    if (tirx::UsesVar(condition, f_use_itervar)) {
+    if (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+            condition,
+            [&](const tirx::Var& var) -> ffi::Expected<ffi::WalkResult> {
+              return f_use_itervar(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                              : ffi::WalkResult::Advance();
+            })
+            .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..6257f010ae03 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,28 @@ 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);
     };
+    bool lhs_uses_itervar = ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+                                lhs_expr,
+                                [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+                                  return f_use_itervar(var.get())
+                                             ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                             : ffi::WalkResult::Advance();
+                                })
+                                .has_value();
+    bool rhs_uses_itervar = ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+                                rhs_expr,
+                                [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+                                  return f_use_itervar(var.get())
+                                             ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                             : ffi::WalkResult::Advance();
+                                })
+                                .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
@@ -1368,7 +1385,14 @@ bool MatchBoundConstraints(PrimExpr pred, ffi::Map<PrimVar, Range>* input_iters,
               } 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,
+                             [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+                               return f_use_itervar(var.get())
+                                          ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                          : ffi::WalkResult::Advance();
+                             })
+                             .has_value()) {
                 lhs_expr = sign ? lhs_expr + part : lhs_expr - part;
               } else {
                 rhs_expr = sign ? rhs_expr - part : rhs_expr + part;
@@ -1430,7 +1454,22 @@ bool IterRangeSanityCheck(const ffi::Map<PrimVar, Range>& iter_ranges) {
   for (const auto& it : iter_ranges) iters.insert(it.first);
   auto f = [&](const VarNode* var) { return iters.count(ffi::GetRef<Var>(var)); };
   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,
+            [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+              return f(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                  : ffi::WalkResult::Advance();
+            })
+            .has_value() ||
+        ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+            it.second->extent,
+            [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+              return f(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                  : ffi::WalkResult::Advance();
+            })
+            .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..3ac7c72e79f2 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>
@@ -237,8 +238,13 @@ class PatternKindAnalyzer : public StmtExprVisitor {
     }
     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,
+              [&vars](const tirx::Var& var) -> ffi::Expected<ffi::WalkResult> {
+                return !vars.count(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                              : ffi::WalkResult::Advance();
+              })
+              .has_value()) {
         return false;
       }
     }
@@ -319,16 +325,19 @@ class PatternKindAnalyzer : public StmtExprVisitor {
   static bool IsPureReducePattern(ffi::Array<tirx::Var> reduce_loops,
                                   ffi::Array<PrimExpr> indices) {
     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, [&](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();
+          });
+      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..475157e291ff 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,14 @@ 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(); })) {
+      if (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+              fn->body,
+              [=](const tirx::Var& var) -> ffi::Expected<ffi::WalkResult> {
+                return var.get() == buffer_var.get()
+                           ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                           : ffi::WalkResult::Advance();
+              })
+              .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..e12aae5c8c6c 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>
 
@@ -74,7 +75,13 @@ 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,
+                [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+                  return f_find(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                           : ffi::WalkResult::Advance();
+                })
+                .has_value()) {
           return false;
         }
       }
diff --git a/src/s_tir/schedule/analysis/analysis.cc b/src/s_tir/schedule/analysis/analysis.cc
index 229c03e1ca04..4f091283a05c 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>
@@ -1834,8 +1835,14 @@ 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_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();
+               })
+               .has_value()) {
         desc_loop = desc_loops[i];
         iter_type_desc = iter_types_desc[i];
         break;
@@ -1869,8 +1876,14 @@ 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_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();
+              })
+              .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..073f95a16ffc 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"
@@ -554,9 +555,14 @@ bool ReductionIterNotIndexOutputBuffer(const SBlock& block) {
   }
 
   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,
+               [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+                 return reduction_block_iters.count(var.get())
+                            ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                            : ffi::WalkResult::Advance();
+               })
+        .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..8206c85f8434 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,14 @@ 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);  //
-      });
+      return ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+                 expr,
+                 [&var_set](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+                   return var_set.count(var.get())
+                              ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                              : ffi::WalkResult::Advance();
+                 })
+          .has_value();
     };
   };
   auto use_outer_loop_vars = make_uses_var(outer_iters);
@@ -336,7 +337,13 @@ Stmt GenerateOuterInit(const Stmt& block_init, const SBlockRealize& inner_realiz
     const IterVar& old_iter_var = inner_block->iter_vars[i];
     const PrimExpr& iter_value = inner_realize->iter_values[i];
     if (old_iter_var->iter_type == IterVarType::kDataPar &&
-        UsesVar(block_init, old_iter_var->var)) {
+        ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+            block_init,
+            [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();
+            })
+            .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);
@@ -359,7 +366,13 @@ Stmt GenerateOuterInit(const Stmt& block_init, const SBlockRealize& inner_realiz
   for (const ForNode* loop : loops) {
     bool is_init_loop = false;
     for (const PrimExpr& init_binding : iter_values) {
-      if (UsesVar(init_binding, loop->loop_var)) {
+      if (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+              init_binding,
+              [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();
+              })
+              .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..b7ad748c4407 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,13 @@ 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; })) {
+    if (!ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+             binding,
+             [v = loop_var.get()](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+               return var.get() == v ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                     : ffi::WalkResult::Advance();
+             })
+             .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..057a2571c28c 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,14 @@ 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;
-    };
-    if (UsesVar(loop->extent, f_contain)) {
+    auto result = ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+        loop->extent, [&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 (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 +1105,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 find_inner_var = [&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, find_inner_var);
+    if (!result.has_value()) {
+      result = ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(copy->extent, find_inner_var);
+    }
+    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..e064f1b46917 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,13 @@ 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; })) {
+        if (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+                binding,
+                [v = loop_var.get()](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+                  return var.get() == v ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                        : ffi::WalkResult::Advance();
+                })
+                .has_value()) {
           const ForNode* loop = TVM_SREF_TO_FOR(loop_sref);
           throw LoopHeightError(mod, ffi::GetRef<For>(loop), ffi::GetRef<SBlock>(block));
         }
@@ -169,7 +176,15 @@ 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;
+  return ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+             pred,
+             [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+               return uses_discarded_loop(var.get())
+                          ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                          : ffi::WalkResult::Advance();
+             }).has_value()
+             ? IntImm::Bool(true)
+             : pred;
 }
 
 StmtSRef DecomposeReduction(ScheduleState self, const StmtSRef& block_sref,
@@ -245,7 +260,13 @@ StmtSRef DecomposeReduction(ScheduleState self, const StmtSRef& block_sref,
     const VarNode* loop_var = loops[i]->StmtAs<ForNode>()->loop_var.get();
     bool discarded = true;
     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,
+               [v = loop_var](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+                 return var.get() == v ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                       : ffi::WalkResult::Advance();
+               })
+               .has_value()) {
         continue;
       }
       // The loop is related to init block bindings;
@@ -897,8 +918,13 @@ class RFactorBlockCreator : public BaseBlockCreator {
     IterVar old_iter = old_block_realize_->block->iter_vars[idx];
     PrimExpr old_binding = old_block_realize_->iter_values[idx];
     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,
+             [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();
+             })
+             .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..a68ad1bf2e36 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,13 @@ 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)) {
+      if (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+              extent,
+              [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+                return is_loop_var(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                              : ffi::WalkResult::Advance();
+              })
+              .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..5297a41d8faf 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,16 @@ 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);
-                              });
+        bool uses_block_var =
+            active_block_vars.size() &&
+            ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+                cond,
+                [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+                  return active_block_vars.count(var.get())
+                             ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                             : ffi::WalkResult::Advance();
+                })
+                .has_value();
 
         std::unordered_set<const VarNode*> let_bindings_used;
 
@@ -389,18 +397,19 @@ 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 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());
-      });
+      bool uses_loop_var =
+          ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+              expr,
+              [&](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());
+                }
+                return matches ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                               : ffi::WalkResult::Advance();
+              })
+              .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..08be04c68405 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,24 @@ 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;
+    if (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+            op->min,
+            [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+              return f_vset_contains(var.get())
+                         ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                         : ffi::WalkResult::Advance();
+            })
+            .has_value() ||
+        ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+            op->extent,
+            [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+              return f_vset_contains(var.get())
+                         ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                         : ffi::WalkResult::Advance();
+            })
+            .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 +317,14 @@ 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(); })) {
+    if (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+            cond,
+            [this](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+              return var.get() == current_var_.get()
+                         ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                         : ffi::WalkResult::Advance();
+            })
+            .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..26ce3e6df0a3 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,24 @@ 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);
+          depends_on_thread_index =
+              depends_on_thread_index &&
+              ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+                  curr_index,
+                  [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+                    return f_uses_thread_index(var.get())
+                               ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                               : ffi::WalkResult::Advance();
+                  })
+                  .has_value() &&
+              ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+                  prev_index,
+                  [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+                    return f_uses_thread_index(var.get())
+                               ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                               : ffi::WalkResult::Advance();
+                  })
+                  .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 9d3475eeebf1..000000000000
--- a/src/tirx/analysis/var_touch.cc
+++ /dev/null
@@ -1,53 +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/ffi/extra/structural_visit.h>
-#include <tvm/tirx/analysis.h>
-
-namespace tvm {
-namespace tirx {
-
-bool UsesVar(const Stmt& stmt, std::function<bool(const VarNode*)> var_set) {
-  auto result = ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-      stmt, [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
-        if (var_set(var.get())) {
-          return ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var));
-        }
-        return ffi::WalkResult::Advance();
-      });
-  return result.has_value();
-}
-
-bool UsesVar(const PrimExpr& expr, std::function<bool(const VarNode*)> var_set) {
-  auto result = ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-      expr, [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
-        if (var_set(var.get())) {
-          return ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var));
-        }
-        return ffi::WalkResult::Advance();
-      });
-  return result.has_value();
-}
-
-}  // namespace tirx
-}  // namespace tvm
diff --git a/src/tirx/script/printer/for_loop.cc b/src/tirx/script/printer/for_loop.cc
index 363a2d466f0b..c93309a1240e 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,14 @@ 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);
-        });
+        return ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+                   e,
+                   [&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();
+                   })
+            .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..cb7e505f1ee9 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,15 @@ 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; });
+      return expr.defined() &&
+             ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+                 expr,
+                 [var](const Var& candidate) -> ffi::Expected<ffi::WalkResult> {
+                   return candidate.get() == var
+                              ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(candidate))
+                              : ffi::WalkResult::Advance();
+                 })
+                 .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..f826eae2cd62 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,14 @@ 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(); }))
+    TVM_FFI_ICHECK(!ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+                        local_index,
+                        [this](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+                          return var.get() == warp_index_.get()
+                                     ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
+                                     : ffi::WalkResult::Advance();
+                        })
+                        .has_value())
         << "LowerWarpMemory failed to rewrite load to shuffle for index " << op->indices[0]
         << " local_index=" << local_index;
 

From 847361ecf7d0192bc6e815c33961154f4bd847e3 Mon Sep 17 00:00:00 2001
From: tqchen <tianqic@nvidia.com>
Date: Thu, 10 Sep 2026 13:48:20 +0000
Subject: [PATCH 4/4] [REFACTOR][TIR] Reuse structural walk callbacks

---
 src/arith/detect_linear_equation.cc           | 24 ++++----
 src/arith/int_set.cc                          | 14 ++---
 src/arith/ir_mutator_with_analyzer.h          | 12 ++--
 src/arith/iter_affine_map.cc                  | 55 ++++++-------------
 src/relax/analysis/tir_op_pattern_kind.cc     | 26 ++++-----
 .../transform/rewrite_dataflow_reshape.cc     | 13 ++---
 .../postproc/rewrite_reduction_block.cc       | 12 ++--
 src/s_tir/schedule/analysis/analysis.cc       | 26 ++++-----
 src/s_tir/schedule/analysis/reducer.cc        | 14 ++---
 .../schedule/primitive/blockize_tensorize.cc  | 39 ++++++-------
 src/s_tir/schedule/primitive/for_kind.cc      | 12 ++--
 .../schedule/primitive/loop_transformation.cc | 17 +++---
 src/s_tir/schedule/primitive/reduction.cc     | 48 +++++++---------
 src/s_tir/transform/compact_buffer_region.cc  | 12 ++--
 src/s_tir/transform/hoist_expression.cc       | 37 ++++++-------
 src/s_tir/transform/loop_partition.cc         | 35 ++++--------
 src/s_tir/transform/thread_storage_sync.cc    | 23 +++-----
 src/tirx/script/printer/for_loop.cc           | 14 ++---
 src/tirx/transform/ir_utils.cc                | 13 ++---
 src/tirx/transform/lower_warp_memory.cc       | 13 ++---
 20 files changed, 180 insertions(+), 279 deletions(-)

diff --git a/src/arith/detect_linear_equation.cc b/src/arith/detect_linear_equation.cc
index 0f97d2c96018..a2a8ffdd5df2 100644
--- a/src/arith/detect_linear_equation.cc
+++ b/src/arith/detect_linear_equation.cc
@@ -110,13 +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 (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-            e,
-            [this](const Var& var) -> ffi::Expected<ffi::WalkResult> {
-              return var.get() == var_.get() ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
-                                             : ffi::WalkResult::Advance();
-            })
-            .has_value()) {
+    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 {
@@ -163,17 +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 (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-            coeff[i - 2],
-            [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
-              return vset_contains(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
-                                              : ffi::WalkResult::Advance();
-            })
-            .has_value()) {
+    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 3aba438f550e..9063c0b6208c 100644
--- a/src/arith/int_set.cc
+++ b/src/arith/int_set.cc
@@ -600,15 +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 (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-              index,
-              [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();
-              })
-              .has_value()) {
+      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 8c40aa24fa4f..d05e8e9a5f86 100644
--- a/src/arith/ir_mutator_with_analyzer.h
+++ b/src/arith/ir_mutator_with_analyzer.h
@@ -109,14 +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 (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-            condition,
-            [&](const tirx::Var& var) -> ffi::Expected<ffi::WalkResult> {
-              return f_use_itervar(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
-                                              : ffi::WalkResult::Advance();
-            })
-            .has_value()) {
+    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 6257f010ae03..534b826d61f0 100644
--- a/src/arith/iter_affine_map.cc
+++ b/src/arith/iter_affine_map.cc
@@ -1349,22 +1349,14 @@ 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);
     };
-    bool lhs_uses_itervar = ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-                                lhs_expr,
-                                [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
-                                  return f_use_itervar(var.get())
-                                             ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
-                                             : ffi::WalkResult::Advance();
-                                })
-                                .has_value();
-    bool rhs_uses_itervar = ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-                                rhs_expr,
-                                [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
-                                  return f_use_itervar(var.get())
-                                             ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
-                                             : ffi::WalkResult::Advance();
-                                })
-                                .has_value();
+    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 (lhs_uses_itervar || rhs_uses_itervar) {
       // At least it uses one input iter
@@ -1378,21 +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 (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-                             part,
-                             [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
-                               return f_use_itervar(var.get())
-                                          ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
-                                          : ffi::WalkResult::Advance();
-                             })
-                             .has_value()) {
+              } 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;
@@ -1453,21 +1438,13 @@ 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 (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-            it.second->min,
-            [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
-              return f(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
-                                  : ffi::WalkResult::Advance();
-            })
-            .has_value() ||
-        ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-            it.second->extent,
-            [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
-              return f(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
-                                  : ffi::WalkResult::Advance();
-            })
-            .has_value()) {
+    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;
     }
   }
diff --git a/src/relax/analysis/tir_op_pattern_kind.cc b/src/relax/analysis/tir_op_pattern_kind.cc
index 3ac7c72e79f2..341d7534cfc2 100644
--- a/src/relax/analysis/tir_op_pattern_kind.cc
+++ b/src/relax/analysis/tir_op_pattern_kind.cc
@@ -236,15 +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 (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-              load_index,
-              [&vars](const tirx::Var& var) -> ffi::Expected<ffi::WalkResult> {
-                return !vars.count(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
-                                              : ffi::WalkResult::Advance();
-              })
-              .has_value()) {
+      if (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(load_index, walkfn).has_value()) {
         return false;
       }
     }
@@ -324,14 +322,14 @@ 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) {
-      auto result = ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-          e, [&](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();
-          });
+      auto result = ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(e, walkfn);
       if (result.has_value()) {
         tirx::Var var = result.value()->value.cast<tirx::Var>();
         int id =
diff --git a/src/relax/transform/rewrite_dataflow_reshape.cc b/src/relax/transform/rewrite_dataflow_reshape.cc
index 475157e291ff..e99f10219d2a 100644
--- a/src/relax/transform/rewrite_dataflow_reshape.cc
+++ b/src/relax/transform/rewrite_dataflow_reshape.cc
@@ -42,14 +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 (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-              fn->body,
-              [=](const tirx::Var& var) -> ffi::Expected<ffi::WalkResult> {
-                return var.get() == buffer_var.get()
-                           ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
-                           : ffi::WalkResult::Advance();
-              })
-              .has_value()) {
+      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 e12aae5c8c6c..9a7a4acf6e7a 100644
--- a/src/s_tir/meta_schedule/postproc/rewrite_reduction_block.cc
+++ b/src/s_tir/meta_schedule/postproc/rewrite_reduction_block.cc
@@ -68,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();
@@ -75,13 +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 (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-                binding,
-                [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
-                  return f_find(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
-                                           : ffi::WalkResult::Advance();
-                })
-                .has_value()) {
+        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 4f091283a05c..20776081f4c5 100644
--- a/src/s_tir/schedule/analysis/analysis.cc
+++ b/src/s_tir/schedule/analysis/analysis.cc
@@ -1827,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];
@@ -1835,14 +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 (!ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-               residual,
-               [&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();
-               })
-               .has_value()) {
+      if (!ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(residual, desc_walkfn).has_value()) {
         desc_loop = desc_loops[i];
         iter_type_desc = iter_types_desc[i];
         break;
@@ -1876,14 +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 (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-              residual,
-              [&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();
-              })
-              .has_value()) {
+      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 073f95a16ffc..c122ea66a091 100644
--- a/src/s_tir/schedule/analysis/reducer.cc
+++ b/src/s_tir/schedule/analysis/reducer.cc
@@ -554,15 +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 ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-               expr,
-               [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
-                 return reduction_block_iters.count(var.get())
-                            ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
-                            : ffi::WalkResult::Advance();
-               })
-        .has_value();
+    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 8206c85f8434..3dea0688496f 100644
--- a/src/s_tir/schedule/primitive/blockize_tensorize.cc
+++ b/src/s_tir/schedule/primitive/blockize_tensorize.cc
@@ -106,14 +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 ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-                 expr,
-                 [&var_set](const Var& var) -> ffi::Expected<ffi::WalkResult> {
-                   return var_set.count(var.get())
-                              ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
-                              : ffi::WalkResult::Advance();
-                 })
-          .has_value();
+      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);
@@ -336,14 +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 &&
-        ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-            block_init,
-            [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();
-            })
-            .has_value()) {
+        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);
@@ -365,14 +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 (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-              init_binding,
-              [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();
-              })
-              .has_value()) {
+      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 b7ad748c4407..b23d50d0f41f 100644
--- a/src/s_tir/schedule/primitive/for_kind.cc
+++ b/src/s_tir/schedule/primitive/for_kind.cc
@@ -98,13 +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 (!ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-             binding,
-             [v = loop_var.get()](const Var& var) -> ffi::Expected<ffi::WalkResult> {
-               return var.get() == v ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
-                                     : ffi::WalkResult::Advance();
-             })
-             .has_value()) {
+    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 057a2571c28c..6c354769df34 100644
--- a/src/s_tir/schedule/primitive/loop_transformation.cc
+++ b/src/s_tir/schedule/primitive/loop_transformation.cc
@@ -908,12 +908,11 @@ StmtSRef Fuse(ScheduleState self, const ffi::Array<StmtSRef>& loop_srefs,
     outer_loop_sref = sref;
     outer_loop = loop;
     CheckLoopStartsWithZero(self, sref, analyzer.get());
-    auto result = ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-        loop->extent, [&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();
-        });
+    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();
+    };
+    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,
@@ -1105,13 +1104,13 @@ For ConstructNewLoopChain(const ScheduleState& self, std::vector<const StmtSRefN
     } else {
       n->body = loop_sref->StmtAs<ForNode>()->body;
     }
-    auto find_inner_var = [&inner_vars](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+    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();
     };
-    auto result = ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(copy->min, find_inner_var);
+    auto result = ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(copy->min, walkfn);
     if (!result.has_value()) {
-      result = ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(copy->extent, find_inner_var);
+      result = ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(copy->extent, walkfn);
     }
     if (result.has_value()) {
       Var used_var = result.value()->value.cast<Var>();
diff --git a/src/s_tir/schedule/primitive/reduction.cc b/src/s_tir/schedule/primitive/reduction.cc
index e064f1b46917..a96face461b2 100644
--- a/src/s_tir/schedule/primitive/reduction.cc
+++ b/src/s_tir/schedule/primitive/reduction.cc
@@ -129,13 +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 (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-                binding,
-                [v = loop_var.get()](const Var& var) -> ffi::Expected<ffi::WalkResult> {
-                  return var.get() == v ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
-                                        : ffi::WalkResult::Advance();
-                })
-                .has_value()) {
+        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));
         }
@@ -176,13 +174,11 @@ PrimExpr RewriteInitPredicate(PrimExpr pred,
   auto uses_discarded_loop = [&discarded_loops](const VarNode* var) {
     return discarded_loops.count(var);
   };
-  return ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-             pred,
-             [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
-               return uses_discarded_loop(var.get())
-                          ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
-                          : ffi::WalkResult::Advance();
-             }).has_value()
+  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;
 }
@@ -259,14 +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 (!ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-               expr,
-               [v = loop_var](const Var& var) -> ffi::Expected<ffi::WalkResult> {
-                 return var.get() == v ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
-                                       : ffi::WalkResult::Advance();
-               })
-               .has_value()) {
+      if (!ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(expr, walkfn).has_value()) {
         continue;
       }
       // The loop is related to init block bindings;
@@ -917,14 +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 ||
-        !ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-             old_binding,
-             [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();
-             })
-             .has_value()) {
+        !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 a68ad1bf2e36..7255317fa358 100644
--- a/src/s_tir/transform/compact_buffer_region.cc
+++ b/src/s_tir/transform/compact_buffer_region.cc
@@ -473,13 +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 (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-              extent,
-              [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
-                return is_loop_var(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
-                                              : ffi::WalkResult::Advance();
-              })
-              .has_value()) {
+      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 5297a41d8faf..40a4b47daf3c 100644
--- a/src/s_tir/transform/hoist_expression.cc
+++ b/src/s_tir/transform/hoist_expression.cc
@@ -218,16 +218,14 @@ class HoistInfoCollector : public StmtExprVisitor {
     if (auto info = FindHoistDestination(cond)) {
       if (!info->reached_sequential_node) {
         // Record whether this conditional uses any block variables.
+        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,
-                [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
-                  return active_block_vars.count(var.get())
-                             ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
-                             : ffi::WalkResult::Advance();
-                })
-                .has_value();
+            ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(cond, walkfn).has_value();
 
         std::unordered_set<const VarNode*> let_bindings_used;
 
@@ -397,19 +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 =
-          ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-              expr,
-              [&](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());
-                }
-                return matches ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
-                               : ffi::WalkResult::Advance();
-              })
-              .has_value();
+      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());
+        }
+        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 08be04c68405..e04a48687e40 100644
--- a/src/s_tir/transform/loop_partition.cc
+++ b/src/s_tir/transform/loop_partition.cc
@@ -249,22 +249,12 @@ class PartitionFinder : public StmtExprVisitor {
 
   void VisitStmt_(const ForNode* op) final {
     auto f_vset_contains = [this](const VarNode* var) { return out_vars_.count(var); };
-    if (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-            op->min,
-            [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
-              return f_vset_contains(var.get())
-                         ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
-                         : ffi::WalkResult::Advance();
-            })
-            .has_value() ||
-        ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-            op->extent,
-            [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
-              return f_vset_contains(var.get())
-                         ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
-                         : ffi::WalkResult::Advance();
-            })
-            .has_value()) {
+    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;
     }
 
@@ -317,14 +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 (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-            cond,
-            [this](const Var& var) -> ffi::Expected<ffi::WalkResult> {
-              return var.get() == current_var_.get()
-                         ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
-                         : ffi::WalkResult::Advance();
-            })
-            .has_value()) {
+    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 26ce3e6df0a3..200ef98ae932 100644
--- a/src/s_tir/transform/thread_storage_sync.cc
+++ b/src/s_tir/transform/thread_storage_sync.cc
@@ -238,24 +238,15 @@ class ThreadSyncPlanner : public StorageAccessVisitor {
           auto f_uses_thread_index = [=](const tvm::tirx::VarNode* parameter) {
             return parameter == thread_index_var;
           };
+          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,
-                  [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
-                    return f_uses_thread_index(var.get())
-                               ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
-                               : ffi::WalkResult::Advance();
-                  })
-                  .has_value() &&
-              ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-                  prev_index,
-                  [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
-                    return f_uses_thread_index(var.get())
-                               ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
-                               : ffi::WalkResult::Advance();
-                  })
-                  .has_value();
+              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/script/printer/for_loop.cc b/src/tirx/script/printer/for_loop.cc
index c93309a1240e..4d0f2016cdec 100644
--- a/src/tirx/script/printer/for_loop.cc
+++ b/src/tirx/script/printer/for_loop.cc
@@ -30,14 +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 ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-                   e,
-                   [&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();
-                   })
-            .has_value();
+        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 cb7e505f1ee9..5633ba59865b 100644
--- a/src/tirx/transform/ir_utils.cc
+++ b/src/tirx/transform/ir_utils.cc
@@ -567,15 +567,12 @@ class IRConvertSSA final : public StmtExprMutator {
     if (buffer.get() == var) return true;
 
     auto uses_var = [var](const PrimExpr& expr) {
+      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,
-                 [var](const Var& candidate) -> ffi::Expected<ffi::WalkResult> {
-                   return candidate.get() == var
-                              ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(candidate))
-                              : ffi::WalkResult::Advance();
-                 })
-                 .has_value();
+             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 f826eae2cd62..4f268b40e096 100644
--- a/src/tirx/transform/lower_warp_memory.cc
+++ b/src/tirx/transform/lower_warp_memory.cc
@@ -385,14 +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(!ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
-                        local_index,
-                        [this](const Var& var) -> ffi::Expected<ffi::WalkResult> {
-                          return var.get() == warp_index_.get()
-                                     ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
-                                     : ffi::WalkResult::Advance();
-                        })
-                        .has_value())
+    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;