[REFACTOR][TIR] Use StructuralWalk in place of PostOrderVisit

include/tvm/tirx/stmt_functor.h

@@ -366,15 +366,6 @@ class TVM_DLL StmtExprMutator : public ExprMutator, public StmtMutator {
   Expr VisitExpr_(const BufferRegionNode* op) override;
 };
 
-/*!
- * \brief Recursively visit a statement or expression in post DFS order, applying fvisit.
- * Each node is guaranteed to be visited only once.
- * \param node The statement or expression to be visited.
- * \param fvisit The visitor function to be applied.
- */
-TVM_DLL void PostOrderVisit(const ffi::ObjectRef& node,
-                            std::function<void(const ffi::ObjectRef&)> fvisit);
-
 /*!
  * \brief Substitute the var specified by vmap.
  * \param stmt The source statement to be substituted
@@ -546,16 +537,6 @@ TVM_DLL Stmt SubstituteWithDataTypeLegalization(
 TVM_DLL PrimExpr SubstituteWithDataTypeLegalization(
     PrimExpr expr, std::function<ffi::Optional<PrimExpr>(const Var&)> vmap);
 
-/*!
- * \brief Recursively visit a statement or expression in pre DFS order, applying fvisit.
- * If fvisit returns false, it won't visit the children of the node.
- * \param stmt_or_expr The statement or expression to be visited.
- * \param fvisit The visitor function to be applied. If fvisit returns false, it won't visit the
- * children of the node
- */
-TVM_DLL void PreOrderVisit(const ffi::ObjectRef& stmt_or_expr,
-                           const std::function<bool(const ffi::ObjectRef&)>& fvisit);
-
 /*!
  * \brief Check if the statement contains the specified node type.
  *

python/tvm/relax/utils.py

@@ -183,7 +183,7 @@ def _visit_expr(e: tirx.Expr):
             if isinstance(e, tvm.ir.Var) and e not in tir_var_map:
                 tir_var_map[e] = tvm.ir.Var(e.name, e.ty)
 
-        tirx.stmt_functor.post_order_visit(expr, _visit_expr)
+        tvm_ffi.structural_walk(expr, (tvm.ir.Var, _visit_expr))
 
     def _convert_te_arg(te_args: Any) -> Any:
         """Helper function used to convert Relax expressions to TE tensor.

python/tvm/s_tir/dlight/analysis/common_analysis.py

@@ -23,7 +23,7 @@
 from collections import namedtuple
 from typing import Literal
 
-from tvm_ffi import get_global_func
+from tvm_ffi import get_global_func, structural_walk
 
 from tvm import ir, s_tir, tirx
 from tvm.s_tir import Schedule
@@ -423,7 +423,7 @@ def _collect_tir_var(expr):
         if ir.is_prim_var(expr):
             tir_vars.add(expr)
 
-    tirx.stmt_functor.post_order_visit(expr, _collect_tir_var)
+    structural_walk(expr, (tirx.Var, _collect_tir_var))
     return tir_vars
 
 

python/tvm/s_tir/dlight/gpu/fallback.py

@@ -17,6 +17,8 @@
 # pylint: disable=missing-docstring
 """A fallback schedule rule for GPU operators."""
 
+import tvm_ffi
+
 from tvm import s_tir, tirx
 from tvm.target import Target
 
@@ -40,7 +42,7 @@ def _visit(node):
         elif isinstance(node, tirx.For) and node.kind == tirx.ForKind.THREAD_BINDING:
             found = True
 
-    tirx.stmt_functor.post_order_visit(stmt, _visit)
+    tvm_ffi.structural_walk(stmt, ((tirx.AttrStmt, tirx.For), _visit))
     return found
 
 

python/tvm/s_tir/dlight/gpu/general_reduction.py

@@ -17,6 +17,8 @@
 # pylint: disable=invalid-name
 """Reduction rule for operators including softmax, layer norm, RMS norm, etc"""
 
+import tvm_ffi
+
 from tvm import arith, ir, s_tir, tirx
 from tvm.target import Target
 
@@ -167,8 +169,8 @@ def _visit_expr(e: tirx.Expr):
                 buffer = buffer_read.buffer
                 if buffer in reduced_buffers:
                     for read_range in buffer_read.region:
-                        tirx.stmt_functor.post_order_visit(read_range.min, _visit_expr)
-                        tirx.stmt_functor.post_order_visit(read_range.extent, _visit_expr)
+                        tvm_ffi.structural_walk(read_range.min, (tirx.Var, _visit_expr))
+                        tvm_ffi.structural_walk(read_range.extent, (tirx.Var, _visit_expr))
 
             s_loops = []
             other_loops = []

python/tvm/tirx/stmt_functor.py

@@ -980,36 +980,6 @@ def visit_expr(self, expr):
         return ExprMutator.visit_expr(self, expr)
 
 
-def post_order_visit(node, fvisit):
-    """Recursively visit a statement or expression in post DFS order, applying fvisit.
-       Each node is guaranteed to be visited only once.
-
-    Parameters
-    ----------
-    node : tvm.tirx.Stmt or tvm.ir.Expr
-        The statement or expression to visit.
-
-    fvisit: function
-        The visitor function.
-    """
-    return _ffi_api.PostOrderVisit(node, fvisit)  # type: ignore
-
-
-def pre_order_visit(node, fvisit):
-    """Recursively visit a statement or expression in pre-order, applying fvisit.
-       If fvisit returns False, it won't visit the children of the node.
-
-    Parameters
-    ----------
-    node : tvm.tirx.Stmt or tvm.ir.Expr
-        The statement or expression to visit.
-
-    fvisit: function of the signature Object -> bool
-        The visitor function.
-    """
-    return _ffi_api.PreOrderVisit(node, fvisit)  # type: ignore
-
-
 def substitute(node, vmap):
     """Substitute the var specified by vmap.
 

src/backend/trn/codegen/codegen_trn.cc

@@ -22,6 +22,7 @@
  */
 #include "codegen_trn.h"
 
+#include <tvm/ffi/extra/structural_visit.h>
 #include <tvm/ir/prim/expr.h>
 #include <tvm/runtime/logging.h>
 #include <tvm/tirx/transform.h>
@@ -309,14 +310,18 @@ std::string CodeGenTrainium::PrintIndices(const Array<PrimExpr>& indices) {
   ctx_.buffer_index = 0;
   ctx_.used_var_cnt = 0;
   for (size_t i = 0; i < indices.size(); ++i) {
-    PreOrderVisit(indices[i], [&](const ffi::ObjectRef& node) {
-      if (const auto* v = node.as<VarNode>()) {
-        if (ctx_.tensorized_loop_vars.count(v)) {
-          ctx_.used_var_cnt++;
-        }
+    std::unordered_set<const ffi::Object*> visited;
+    auto walk_fn = [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+      const VarNode* v = var.get();
+      if (!visited.insert(v).second) {
+        return ffi::WalkResult::Advance();
       }
-      return true;
-    });
+      if (ctx_.tensorized_loop_vars.count(v)) {
+        ctx_.used_var_cnt++;
+      }
+      return ffi::WalkResult::Advance();
+    };
+    ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(indices[i], walk_fn);
   }
   for (size_t i = 0; i < indices.size(); ++i) {
     if (i != 0) {
@@ -515,22 +520,22 @@ void CodeGenTrainium::VisitExpr_(const CallNode* op, std::ostream& os) {  // NOL
     LOG(FATAL) << "Trainium codegen does not support call to " << op->op;
   }
   if (ctx_.mask.defined()) {
-    PreOrderVisit(ctx_.mask, [&](const ffi::ObjectRef& node) {
-      if (const auto* v = node.as<VarNode>()) {
-        if (ctx_.tensorized_loop_vars.count(v)) {
-          TVM_FFI_ICHECK(ctx_.loopvar2dim.count(v))
-              << "nki_dim must be specified for tensorized loop variables used in mask. However, "
-                 "it is not specified for "
-              << ffi::GetRef<Var>(v);
-          auto dim_str = ctx_.loopvar2dim[v];
-          TVM_FFI_ICHECK(dim_str == "P" || dim_str == "F")
-              << "Only nki_dim = P or F is allowed for tensorized loop variables used in mask. "
-                 "However, "
-              << ffi::GetRef<Var>(v) << " has nki_dim = " << dim_str;
-        }
+    auto walk_fn = [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+      const VarNode* v = var.get();
+      if (ctx_.tensorized_loop_vars.count(v)) {
+        TVM_FFI_ICHECK(ctx_.loopvar2dim.count(v))
+            << "nki_dim must be specified for tensorized loop variables used in mask. However, "
+               "it is not specified for "
+            << ffi::GetRef<Var>(v);
+        auto dim_str = ctx_.loopvar2dim[v];
+        TVM_FFI_ICHECK(dim_str == "P" || dim_str == "F")
+            << "Only nki_dim = P or F is allowed for tensorized loop variables used in mask. "
+               "However, "
+            << ffi::GetRef<Var>(v) << " has nki_dim = " << dim_str;
       }
-      return true;
-    });
+      return ffi::WalkResult::Advance();
+    };
+    ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(ctx_.mask, walk_fn);
     os << ", mask=" << PrintExpr(ctx_.mask);
   }
   os << ")";

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>
@@ -260,15 +261,14 @@ class PatternKindAnalyzer : public StmtExprVisitor {
         return false;
       }
     }
+    auto walk_fn = [&](const tirx::Var& var) -> ffi::Expected<ffi::WalkResult> {
+      if (auto prim_var = var.as<tirx::PrimVar>()) {
+        vars.erase(prim_var.value().get());
+      }
+      return ffi::WalkResult::Advance();
+    };
     for (const PrimExpr& index : load->indices) {
-      PreOrderVisit(index, [&](const ffi::ObjectRef& node) {
-        if (auto var = node.as<tirx::PrimVar>()) {
-          if (vars.count(var.value().get())) {
-            vars.erase(var.value().get());
-          }
-        }
-        return true;
-      });
+      ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(index, walk_fn);
     }
     return !vars.empty();
   }

src/relax/script/printer/dependent_type.cc

@@ -17,6 +17,7 @@
  * under the License.
  */
 #include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
 #include <tvm/tirx/stmt_functor.h>
 
 #include "./utils.h"
@@ -44,13 +45,15 @@ ExprDoc PrintShapeVar(const PrimExpr& e, const AccessPath& e_p, const IRDocsifie
   // Step 2. Figure out if the PrimExpr contains at least a func var
   bool func_var_mode = false;
   if (f != nullptr) {
-    tirx::PostOrderVisit(e, [f, &func_var_mode](const ffi::ObjectRef& obj) -> void {
-      if (auto var = obj.as<tirx::PrimVar>()) {
-        if (f->func_vars->count(var.value().get())) {
+    auto walk_fn = [f, &func_var_mode](const tirx::Var& var) -> ffi::Expected<ffi::WalkResult> {
+      if (auto prim_var = var.as<tirx::PrimVar>()) {
+        if (f->func_vars->count(prim_var.value().get())) {
           func_var_mode = true;
         }
       }
-    });
+      return ffi::WalkResult::Advance();
+    };
+    ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(e, walk_fn);
   }
   // Step 3. Stringify the PrimExpr if func var exists
   bool is_bare_type_var = false;

src/relax/transform/rewrite_cuda_graph.cc

@@ -50,6 +50,7 @@
  * with `CUDAGraphRewriter`.
  */
 #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/relax/analysis.h>
@@ -483,17 +484,22 @@ class CUDAGraphRewritePlanner : public ExprVisitor {
                 [[maybe_unused]] std::vector<const VarNode*>* vars_collector = nullptr,
                 std::vector<tirx::PrimVar>* tir_vars_collector = nullptr) {
     bool is_static = true;
-    tirx::PostOrderVisit(expr, [&](const ffi::ObjectRef& e) {
-      if (auto var = e.as<tirx::PrimVar>()) {
-        if (!capture_symbolic_vars_.count(var.value())) {
-          is_static = false;
-          return;
-        }
-        if (tir_vars_collector != nullptr) {
-          tir_vars_collector->push_back(var.value());
-        }
+    std::unordered_set<const ffi::Object*> visited;
+    auto walk_fn = [&](const tirx::Var& var) -> ffi::Expected<ffi::WalkResult> {
+      auto prim_var = var.as<tirx::PrimVar>();
+      if (!prim_var || !visited.insert(prim_var.value().get()).second) {
+        return ffi::WalkResult::Advance();
+      }
+      if (!capture_symbolic_vars_.count(prim_var.value())) {
+        is_static = false;
+        return ffi::WalkResult::Advance();
       }
-    });
+      if (tir_vars_collector != nullptr) {
+        tir_vars_collector->push_back(prim_var.value());
+      }
+      return ffi::WalkResult::Advance();
+    };
+    ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(expr, walk_fn);
     return is_static;
   }
 

src/s_tir/analysis/sblock_buffer_access_lca_detector.cc

@@ -23,6 +23,7 @@
  */
 
 #include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/tirx/analysis.h>
 #include <tvm/tirx/stmt_functor.h>
@@ -150,12 +151,12 @@ class LCADetector : public StmtExprVisitor {
     auto do_collect_itervar_scope = [this](const IterVar& itervar,
                                            const PrimExpr& binding) -> const ScopeInfo* {
       const ScopeInfo* highest_scope = nullptr;
-      PostOrderVisit(binding, [this, &highest_scope](const ffi::ObjectRef& obj) {
-        if (auto var = obj.as<PrimVar>()) {
-          const VarNode* loop_var = var.value().get();
+      auto walk_fn = [this, &highest_scope](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+        if (auto prim_var = var.as<PrimVar>()) {
+          const VarNode* loop_var = prim_var.value().get();
           auto it = loop_scope_map_.find(loop_var);
           if (it == loop_scope_map_.end()) {
-            return;
+            return ffi::WalkResult::Advance();
           }
           const ScopeInfo* scope = it->second->parent_scope_info;
           if (highest_scope == nullptr) {
@@ -164,7 +165,9 @@ class LCADetector : public StmtExprVisitor {
             highest_scope = scope;
           }
         }
-      });
+        return ffi::WalkResult::Advance();
+      };
+      ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(binding, walk_fn);
       return highest_scope;
     };
 
@@ -200,25 +203,28 @@ class LCADetector : public StmtExprVisitor {
       const BufferVar& buffer = region->buffer;
       const ScopeInfo* scope = ancestor_scopes_.back();
 
-      auto handle_itervar = [&opaque_var_scope, &scope](const ffi::ObjectRef& obj) {
-        if (auto var = obj.as<PrimVar>()) {
-          const VarNode* iter_var = var.value().get();
+      auto handle_itervar = [&opaque_var_scope,
+                             &scope](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+        if (auto prim_var = var.as<PrimVar>()) {
+          const VarNode* iter_var = prim_var.value().get();
           auto dom_scope_it = opaque_var_scope.find(iter_var);
           if (dom_scope_it == opaque_var_scope.end()) {
-            return;
+            return ffi::WalkResult::Advance();
           }
           // find the highest loop scope the accessed buffer index has
           // loop carried dependencies to (via opaque iter var binding).
           if (dom_scope_it->second->depth < scope->depth) {
             scope = dom_scope_it->second;
           }
         }
+        return ffi::WalkResult::Advance();
       };
 
       // visit region min and max to find the lowest legal lca scope
       for (const Range& range : region->region) {
-        PostOrderVisit(range->min, handle_itervar);
-        PostOrderVisit(range->min + range->extent - 1, handle_itervar);
+        ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(range->min, handle_itervar);
+        ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(range->min + range->extent - 1,
+                                                        handle_itervar);
       }
 
       // the scope should be above `highest_reduce_scope` for reduce output buffer.