diff --git a/include/tvm/tirx/stmt_functor.h b/include/tvm/tirx/stmt_functor.h index 54cf329a0a43..9aa466f603d7 100644 --- a/include/tvm/tirx/stmt_functor.h +++ b/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 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(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& fvisit); - /*! * \brief Check if the statement contains the specified node type. * diff --git a/python/tvm/relax/utils.py b/python/tvm/relax/utils.py index 78242bd07027..2d9e42523f6b 100644 --- a/python/tvm/relax/utils.py +++ b/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. diff --git a/python/tvm/s_tir/dlight/analysis/common_analysis.py b/python/tvm/s_tir/dlight/analysis/common_analysis.py index 9c29d49d7947..1a697ef7efe9 100644 --- a/python/tvm/s_tir/dlight/analysis/common_analysis.py +++ b/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 diff --git a/python/tvm/s_tir/dlight/gpu/fallback.py b/python/tvm/s_tir/dlight/gpu/fallback.py index 0acb39a6ba0a..0c655666324d 100644 --- a/python/tvm/s_tir/dlight/gpu/fallback.py +++ b/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 diff --git a/python/tvm/s_tir/dlight/gpu/general_reduction.py b/python/tvm/s_tir/dlight/gpu/general_reduction.py index d3d758afc9cb..b79d644ac1fb 100644 --- a/python/tvm/s_tir/dlight/gpu/general_reduction.py +++ b/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 = [] diff --git a/python/tvm/tirx/stmt_functor.py b/python/tvm/tirx/stmt_functor.py index eec1c184f9fc..439e04930b8d 100644 --- a/python/tvm/tirx/stmt_functor.py +++ b/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. diff --git a/src/backend/trn/codegen/codegen_trn.cc b/src/backend/trn/codegen/codegen_trn.cc index da0f6f744b43..5646054de565 100644 --- a/src/backend/trn/codegen/codegen_trn.cc +++ b/src/backend/trn/codegen/codegen_trn.cc @@ -22,6 +22,7 @@ */ #include "codegen_trn.h" +#include #include #include #include @@ -309,14 +310,18 @@ std::string CodeGenTrainium::PrintIndices(const Array& 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()) { - if (ctx_.tensorized_loop_vars.count(v)) { - ctx_.used_var_cnt++; - } + std::unordered_set visited; + auto walk_fn = [&](const Var& var) -> ffi::Expected { + 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(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()) { - 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(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(v) << " has nki_dim = " << dim_str; - } + auto walk_fn = [&](const Var& var) -> ffi::Expected { + 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(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(v) << " has nki_dim = " << dim_str; } - return true; - }); + return ffi::WalkResult::Advance(); + }; + ffi::StructuralWalk(ctx_.mask, walk_fn); os << ", mask=" << PrintExpr(ctx_.mask); } os << ")"; diff --git a/src/relax/analysis/tir_op_pattern_kind.cc b/src/relax/analysis/tir_op_pattern_kind.cc index 5287e56d96dc..de8a8d06d916 100644 --- a/src/relax/analysis/tir_op_pattern_kind.cc +++ b/src/relax/analysis/tir_op_pattern_kind.cc @@ -19,6 +19,7 @@ #include #include +#include #include #include #include @@ -260,15 +261,14 @@ class PatternKindAnalyzer : public StmtExprVisitor { return false; } } + auto walk_fn = [&](const tirx::Var& var) -> ffi::Expected { + if (auto prim_var = var.as()) { + 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()) { - if (vars.count(var.value().get())) { - vars.erase(var.value().get()); - } - } - return true; - }); + ffi::StructuralWalk(index, walk_fn); } return !vars.empty(); } diff --git a/src/relax/script/printer/dependent_type.cc b/src/relax/script/printer/dependent_type.cc index af1f126fceb4..606e37a63f1b 100644 --- a/src/relax/script/printer/dependent_type.cc +++ b/src/relax/script/printer/dependent_type.cc @@ -17,6 +17,7 @@ * under the License. */ #include +#include #include #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()) { - if (f->func_vars->count(var.value().get())) { + auto walk_fn = [f, &func_var_mode](const tirx::Var& var) -> ffi::Expected { + if (auto prim_var = var.as()) { + if (f->func_vars->count(prim_var.value().get())) { func_var_mode = true; } } - }); + return ffi::WalkResult::Advance(); + }; + ffi::StructuralWalk(e, walk_fn); } // Step 3. Stringify the PrimExpr if func var exists bool is_bare_type_var = false; diff --git a/src/relax/transform/rewrite_cuda_graph.cc b/src/relax/transform/rewrite_cuda_graph.cc index 8515d9df3e63..99b0f2d2d60a 100644 --- a/src/relax/transform/rewrite_cuda_graph.cc +++ b/src/relax/transform/rewrite_cuda_graph.cc @@ -50,6 +50,7 @@ * with `CUDAGraphRewriter`. */ #include +#include #include #include #include @@ -483,17 +484,22 @@ class CUDAGraphRewritePlanner : public ExprVisitor { [[maybe_unused]] std::vector* vars_collector = nullptr, std::vector* tir_vars_collector = nullptr) { bool is_static = true; - tirx::PostOrderVisit(expr, [&](const ffi::ObjectRef& e) { - if (auto var = e.as()) { - 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 visited; + auto walk_fn = [&](const tirx::Var& var) -> ffi::Expected { + auto prim_var = var.as(); + 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(expr, walk_fn); return is_static; } diff --git a/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc b/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc index 1a2fac941571..2ccacf0ca106 100644 --- a/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc +++ b/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc @@ -23,6 +23,7 @@ */ #include +#include #include #include #include @@ -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()) { - const VarNode* loop_var = var.value().get(); + auto walk_fn = [this, &highest_scope](const Var& var) -> ffi::Expected { + if (auto prim_var = var.as()) { + 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(binding, walk_fn); return highest_scope; }; @@ -200,12 +203,13 @@ 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()) { - const VarNode* iter_var = var.value().get(); + auto handle_itervar = [&opaque_var_scope, + &scope](const Var& var) -> ffi::Expected { + if (auto prim_var = var.as()) { + 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). @@ -213,12 +217,14 @@ class LCADetector : public StmtExprVisitor { 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(range->min, handle_itervar); + ffi::StructuralWalk(range->min + range->extent - 1, + handle_itervar); } // the scope should be above `highest_reduce_scope` for reduce output buffer. diff --git a/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc b/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc index 6a20f4aa4b04..5fe771d5db56 100644 --- a/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc +++ b/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc @@ -17,6 +17,7 @@ * under the License. */ #include +#include #include #include #include @@ -262,13 +263,11 @@ Pass SimplifyForFeatureExtraction() { private: static bool HasBufferLoad(const PrimExpr& expr) { - bool found = false; - PostOrderVisit(expr, [&found](const ffi::ObjectRef& node) { - if (node->IsInstance()) { - found = true; - } - }); - return found; + auto walk_fn = [](const TensorLoad&) -> ffi::Expected { + return ffi::WalkResult::Interrupt(ffi::VisitInterrupt(true)); + }; + auto result = ffi::StructuralWalk(expr, walk_fn); + return result.has_value() ? result.value()->value.cast() : false; } Expr VisitExpr_(const SelectNode* node) final { @@ -798,28 +797,28 @@ void Feature::Init(const BufferStoreNode* store, int n_loops) { info.access_type = AccessType::kWrite; info.multi_indices.push_back({store->indices.begin(), store->indices.end()}); } - PostOrderVisit(store->value, [&buffer_info](const ffi::ObjectRef& obj) -> void { - if (const TensorLoadNode* load = obj.as()) { - BufferVar buffer = load->source.as_or_throw(); - Info& info = buffer_info[buffer]; - switch (info.access_type) { - case AccessType::kRead: - break; - case AccessType::kWrite: - info.access_type = AccessType::kReadWrite; - break; - case AccessType::kReadWrite: - break; - case AccessType::kUnknownRW: - default: - info.access_type = AccessType::kRead; - break; - } - if (info.access_type != AccessType::kReadWrite) { - info.multi_indices.push_back({load->indices.begin(), load->indices.end()}); - } + auto walk_fn = [&buffer_info](const TensorLoad& load) -> ffi::Expected { + BufferVar buffer = load->source.as_or_throw(); + Info& info = buffer_info[buffer]; + switch (info.access_type) { + case AccessType::kRead: + break; + case AccessType::kWrite: + info.access_type = AccessType::kReadWrite; + break; + case AccessType::kReadWrite: + break; + case AccessType::kUnknownRW: + default: + info.access_type = AccessType::kRead; + break; } - }); + if (info.access_type != AccessType::kReadWrite) { + info.multi_indices.push_back({load->indices.begin(), load->indices.end()}); + } + return ffi::WalkResult::Advance(); + }; + ffi::StructuralWalk(store->value, walk_fn); this->sub_features.reserve(buffer_info.size()); for (const auto& kv : buffer_info) { this->sub_features.emplace_back(kv.first, kv.second.access_type, @@ -920,13 +919,15 @@ void Feature::SubFeature::SetReuse(const LoopNest& loop_nest, int64_t top_loop_t BufferVar buffer = this->buffer; // Step 3.1. Collect all `Var`s that appears in the buffer region std::unordered_set region_vars; + auto walk_fn = [®ion_vars](const Var& var) -> ffi::Expected { + if (auto prim_var = var.as()) { + region_vars.insert(prim_var.value().get()); + } + return ffi::WalkResult::Advance(); + }; for (const MultiIndex& multi_index : this->multi_indices) { for (const PrimExpr& index : multi_index) { - PostOrderVisit(index, [®ion_vars](const ffi::ObjectRef& obj) -> void { - if (auto var = obj.as()) { - region_vars.insert(var.value().get()); - } - }); + ffi::StructuralWalk(index, walk_fn); } } // Default case: no reuse diff --git a/src/s_tir/meta_schedule/postproc/rewrite_cooperative_fetch.cc b/src/s_tir/meta_schedule/postproc/rewrite_cooperative_fetch.cc index ad5bb883b3ac..2cea200f41be 100644 --- a/src/s_tir/meta_schedule/postproc/rewrite_cooperative_fetch.cc +++ b/src/s_tir/meta_schedule/postproc/rewrite_cooperative_fetch.cc @@ -16,6 +16,7 @@ * specific language governing permissions and limitations * under the License. */ +#include #include #include #include @@ -91,23 +92,29 @@ bool ParseWarpExecutionAnn(const Schedule& sch, const Instruction& inst) { size_t GetMaxUsedDtypeBytes(SBlock block) { size_t max_bytes = 1; - - tirx::PostOrderVisit(block->body, [&](const ffi::ObjectRef& obj) { - if (const auto* store = obj.as()) { - max_bytes = std::max(max_bytes, store->value.ty().StorageBytes()); - } else if (const auto* load = obj.as()) { - max_bytes = std::max(max_bytes, load->ty.as_or_throw().StorageBytes()); - } else if (const auto* call = obj.as()) { - static const Op& q_multiply_shift_per_axis_op = Op::Get("tirx.q_multiply_shift_per_axis"); - static const Op& q_multiply_shift_op = Op::Get("tirx.q_multiply_shift"); - if (call->op.same_as(q_multiply_shift_per_axis_op) || call->op.same_as(q_multiply_shift_op)) { - // q_multiply_shift uses 64 bit multiply - max_bytes = std::max(max_bytes, 8); - } - } else if (const auto* cast = obj.as()) { - max_bytes = std::max(max_bytes, cast->ty.as_or_throw().StorageBytes()); + auto visit_store = [&](const tirx::BufferStore& store) -> ffi::Expected { + max_bytes = std::max(max_bytes, store->value.ty().StorageBytes()); + return ffi::WalkResult::Advance(); + }; + auto visit_load = [&](const TensorLoad& load) -> ffi::Expected { + max_bytes = std::max(max_bytes, load->ty.as_or_throw().StorageBytes()); + return ffi::WalkResult::Advance(); + }; + auto visit_call = [&](const Call& call) -> ffi::Expected { + static const Op& q_multiply_shift_per_axis_op = Op::Get("tirx.q_multiply_shift_per_axis"); + static const Op& q_multiply_shift_op = Op::Get("tirx.q_multiply_shift"); + if (call->op.same_as(q_multiply_shift_per_axis_op) || call->op.same_as(q_multiply_shift_op)) { + // q_multiply_shift uses 64 bit multiply + max_bytes = std::max(max_bytes, 8); } - }); + return ffi::WalkResult::Advance(); + }; + auto visit_cast = [&](const prim::Cast& cast) -> ffi::Expected { + max_bytes = std::max(max_bytes, cast->ty.as_or_throw().StorageBytes()); + return ffi::WalkResult::Advance(); + }; + ffi::StructuralWalk(block->body, visit_store, visit_load, visit_call, + visit_cast); return max_bytes; } diff --git a/src/s_tir/meta_schedule/postproc/rewrite_tensorize.cc b/src/s_tir/meta_schedule/postproc/rewrite_tensorize.cc index c1b94094069e..638706b0b432 100644 --- a/src/s_tir/meta_schedule/postproc/rewrite_tensorize.cc +++ b/src/s_tir/meta_schedule/postproc/rewrite_tensorize.cc @@ -16,6 +16,7 @@ * specific language governing permissions and limitations * under the License. */ +#include #include #include #include @@ -37,32 +38,32 @@ void CollectTensorizationJobs( const s_tir::Schedule& sch, const ffi::String& func_name, const tirx::PrimFuncNode* func, bool vectorize_init_loop, std::vector>>* jobs) { - tirx::PostOrderVisit(func->body, [=, &jobs](const ffi::ObjectRef& obj) { - if (const auto* block = obj.as()) { - tirx::StmtSRef block_sref = sch->GetSRef(block); - std::string block_name = block_sref->StmtAs()->name_hint; - if (ffi::Optional intrin_name = - s_tir::GetAnn(block_sref, s_tir::attr::meta_schedule_auto_tensorize)) { - if (intrin_name.value() != "") { - jobs->emplace_back(block_name, func_name, [sch, intrin_name](s_tir::SBlockRV block) { - try { - sch->Tensorize(block, intrin_name.value()); - } catch (const std::exception& e) { - LOG(WARNING) << "Tensorize failed with error " << e.what(); - } - }); - } else if (block_name.find("init") && vectorize_init_loop) { - jobs->emplace_back(block_name, func_name, [sch](s_tir::SBlockRV block) { - ffi::Array child_blocks = sch->GetChildBlocks(block); - TVM_FFI_ICHECK(child_blocks.size() == 1); - ffi::Array init_loops = sch->GetLoops(child_blocks[0]); - TVM_FFI_ICHECK(init_loops.size() == 1); - sch->Vectorize(init_loops[0]); - }); - } + auto walk_fn = [=, &jobs](const tirx::SBlock& block) -> ffi::Expected { + tirx::StmtSRef block_sref = sch->GetSRef(block.get()); + std::string block_name = block_sref->StmtAs()->name_hint; + if (ffi::Optional intrin_name = + s_tir::GetAnn(block_sref, s_tir::attr::meta_schedule_auto_tensorize)) { + if (intrin_name.value() != "") { + jobs->emplace_back(block_name, func_name, [sch, intrin_name](s_tir::SBlockRV block) { + try { + sch->Tensorize(block, intrin_name.value()); + } catch (const std::exception& e) { + LOG(WARNING) << "Tensorize failed with error " << e.what(); + } + }); + } else if (block_name.find("init") && vectorize_init_loop) { + jobs->emplace_back(block_name, func_name, [sch](s_tir::SBlockRV block) { + ffi::Array child_blocks = sch->GetChildBlocks(block); + TVM_FFI_ICHECK(child_blocks.size() == 1); + ffi::Array init_loops = sch->GetLoops(child_blocks[0]); + TVM_FFI_ICHECK(init_loops.size() == 1); + sch->Vectorize(init_loops[0]); + }); } } - }); + return ffi::WalkResult::Advance(); + }; + ffi::StructuralWalk(func->body, walk_fn); } class RewriteTensorizeNode : public PostprocNode { diff --git a/src/s_tir/schedule/analysis/analysis.cc b/src/s_tir/schedule/analysis/analysis.cc index 229c03e1ca04..342632b6ceba 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 +#include #include #include #include @@ -143,11 +144,13 @@ ScopeBlockLoopInfo GetScopeBlockLoopInfo(const SBlock& scope_block) { } else { vars = &result.non_spatial_vars; } - PostOrderVisit(iter_value, [vars](const ffi::ObjectRef& obj) { - if (auto var = obj.as()) { - vars->insert(var.value().get()); + auto walk_fn = [vars](const Var& var) -> ffi::Expected { + if (auto prim_var = var.as()) { + vars->insert(prim_var.value().get()); } - }); + return ffi::WalkResult::Advance(); + }; + ffi::StructuralWalk(iter_value, walk_fn); } } @@ -903,36 +906,36 @@ IterVarType GetLoopIterType(const StmtSRef& loop_sref) { int n_spatial = 0; int n_reduce = 0; int n_other = 0; - auto f_visit = [&loop_var, &n_spatial, &n_reduce, &n_other](const ffi::ObjectRef& obj) -> bool { - if (const auto* realize = obj.as()) { - const SBlockNode* block = realize->block.get(); - // Number of block vars and their bindings - TVM_FFI_ICHECK_EQ(realize->iter_values.size(), block->iter_vars.size()); - size_t n = realize->iter_values.size(); - for (size_t i = 0; i < n; ++i) { - const IterVar& iter_var = block->iter_vars[i]; - const PrimExpr& binding = realize->iter_values[i]; - // Categorize the current block var - int* ref = nullptr; - if (iter_var->iter_type == IterVarType::kDataPar) { - ref = &n_spatial; - } else if (iter_var->iter_type == IterVarType::kCommReduce) { - ref = &n_reduce; - } else { - ref = &n_other; - } - // Visit the binding to see if `loop_var` appears - PostOrderVisit(binding, [&ref, &loop_var](const ffi::ObjectRef& obj) -> void { - if (obj.same_as(loop_var)) { - (*ref) += 1; - } - }); + auto f_visit = [&loop_var, &n_spatial, &n_reduce, + &n_other](const SBlockRealize& realize) -> ffi::Expected { + const SBlockNode* block = realize->block.get(); + // Number of block vars and their bindings + TVM_FFI_ICHECK_EQ(realize->iter_values.size(), block->iter_vars.size()); + size_t n = realize->iter_values.size(); + for (size_t i = 0; i < n; ++i) { + const IterVar& iter_var = block->iter_vars[i]; + const PrimExpr& binding = realize->iter_values[i]; + // Categorize the current block var + int* ref = nullptr; + if (iter_var->iter_type == IterVarType::kDataPar) { + ref = &n_spatial; + } else if (iter_var->iter_type == IterVarType::kCommReduce) { + ref = &n_reduce; + } else { + ref = &n_other; } - return false; + // Visit the binding to see if `loop_var` appears + auto walk_fn = [&ref, &loop_var](const Var& var) -> ffi::Expected { + if (var.same_as(loop_var)) { + (*ref) += 1; + } + return ffi::WalkResult::Advance(); + }; + ffi::StructuralWalk(binding, walk_fn); } - return true; + return ffi::WalkResult::Skip(); }; - PreOrderVisit(loop->body, f_visit); + ffi::StructuralWalk(loop->body, f_visit); if (n_other) { return IterVarType::kOpaque; } else if (n_spatial && n_reduce) { @@ -1332,47 +1335,39 @@ bool HasOp(const Stmt& stmt, const ffi::Array& ops) { for (const Op& op : ops) { op_set.insert(op.operator->()); } - bool found = false; - PreOrderVisit(stmt, [&found, &op_set](const ffi::ObjectRef& obj) -> bool { - if (found) { - return false; - } - if (const auto* call = obj.as()) { - if (op_set.count(call->op.operator->())) { - found = true; - } + auto walk_fn = [&op_set](const Call& call) -> ffi::Expected { + if (op_set.count(call->op.operator->())) { + return ffi::WalkResult::Interrupt(ffi::VisitInterrupt(true)); } - return !found; - }); - return found; + return ffi::WalkResult::Advance(); + }; + auto result = ffi::StructuralWalk(stmt, walk_fn); + return result.has_value() ? result.value()->value.cast() : false; } bool HasIfThenElse(const Stmt& stmt) { - bool has_branch = false; - auto f_visit = [&has_branch](const ffi::ObjectRef& obj) -> bool { - if (has_branch) { - // stop visiting - return false; + auto visit_realize = [](const SBlockRealize& realize) -> ffi::Expected { + if (!is_one(realize->predicate)) { + return ffi::WalkResult::Interrupt(ffi::VisitInterrupt(true)); } - if (const auto* realize = obj.as()) { - // Case 1: BlockRealize - if (!is_one(realize->predicate)) { - has_branch = true; - } - } else if (obj->IsInstance() || obj->IsInstance()) { - // Case 2: IfThenElse / Select - has_branch = true; - } else if (const auto* call = obj.as()) { - // Case 3: Call the `if_then_else` operator - static const Op& if_then_else_op = Op::Get("ir.prim.if_then_else"); - if (call->op.same_as(if_then_else_op)) { - has_branch = true; - } + return ffi::WalkResult::Advance(); + }; + auto visit_branch = [](const IfThenElse&) -> ffi::Expected { + return ffi::WalkResult::Interrupt(ffi::VisitInterrupt(true)); + }; + auto visit_select = [](const Select&) -> ffi::Expected { + return ffi::WalkResult::Interrupt(ffi::VisitInterrupt(true)); + }; + auto visit_call = [](const Call& call) -> ffi::Expected { + static const Op& if_then_else_op = Op::Get("ir.prim.if_then_else"); + if (call->op.same_as(if_then_else_op)) { + return ffi::WalkResult::Interrupt(ffi::VisitInterrupt(true)); } - return !has_branch; + return ffi::WalkResult::Advance(); }; - PreOrderVisit(stmt, f_visit); - return has_branch; + auto result = ffi::StructuralWalk(stmt, visit_realize, visit_branch, + visit_select, visit_call); + return result.has_value() ? result.value()->value.cast() : false; } std::tuplebody, [&result](const ffi::ObjectRef& obj) { - if (result == false) { - return false; - } - if (const auto* block = obj.as()) { - for (const IterVar& iter_var : block->iter_vars) { - if (iter_var->iter_type != IterVarType::kDataPar) { - result = false; - return false; - } + auto walk_fn = [](const SBlock& block) -> ffi::Expected { + for (const IterVar& iter_var : block->iter_vars) { + if (iter_var->iter_type != IterVarType::kDataPar) { + return ffi::WalkResult::Interrupt(ffi::VisitInterrupt(false)); } } - return true; - }); - return result; + return ffi::WalkResult::Advance(); + }; + auto result = ffi::StructuralWalk(func->body, walk_fn); + return result.has_value() ? result.value()->value.cast() : true; } std::pair GetCumulativeSpaceAndReductionLength(const s_tir::ScheduleState& self, @@ -1738,23 +1727,20 @@ TensorIntrinDescInfo ExtractTensorIntrinDescInfo(arith::AnalyzerObj* analyzer, const auto* desc_scope_realize = desc_func->body.as(); TVM_FFI_ICHECK(desc_scope_realize); { - auto f_visit = [&](const ffi::ObjectRef& obj) -> bool { - // Extract the block - if (const auto* block = obj.as()) { - info.desc_block = block; - return false; - } - // Extract the loops - if (const auto* loop = obj.as()) { - info.desc_loops.push_back(loop); - info.desc_loop_vars.insert(loop->loop_var.get()); - if (!analyzer->CanProve(loop->min == 0)) { - return false; - } + auto visit_block = [&](const SBlockRealize& block) -> ffi::Expected { + info.desc_block = block.get(); + return ffi::WalkResult::Advance(); + }; + auto visit_loop = [&](const For& loop) -> ffi::Expected { + info.desc_loops.push_back(loop.get()); + info.desc_loop_vars.insert(loop->loop_var.get()); + if (!analyzer->CanProve(loop->min == 0)) { + return ffi::WalkResult::Advance(); } - return true; + return ffi::WalkResult::Advance(); }; - tirx::PostOrderVisit(desc_scope_realize->block->body, f_visit); + ffi::StructuralWalk(desc_scope_realize->block->body, visit_block, + visit_loop); std::reverse(info.desc_loops.begin(), info.desc_loops.end()); TVM_FFI_ICHECK(info.desc_block); } @@ -2016,13 +2002,14 @@ class AutoTensorizeMappingProposer { auto lhs_buffer_it = extractor_->rhs_buffer_map_.find(rhs_buffer); TVM_FFI_ICHECK(lhs_buffer_it != extractor_->rhs_buffer_map_.end()); const BufferVar& lhs_buffer = lhs_buffer_it->second; + auto walk_fn = [&](const Var& var) -> ffi::Expected { + if (auto prim_var = var.as()) { + update_mask(prim_var.value().get(), &lhs_buffer_masks, lhs_buffer_index.at(lhs_buffer)); + } + return ffi::WalkResult::Advance(); + }; for (const PrimExpr& index : extractor_->lhs_buffer_indices_map_.at(lhs_buffer)) { - PreOrderVisit(index, [&](const ffi::ObjectRef& obj) -> bool { - if (auto var = obj.as()) { - update_mask(var.value().get(), &lhs_buffer_masks, lhs_buffer_index.at(lhs_buffer)); - } - return true; - }); + ffi::StructuralWalk(index, walk_fn); } } diff --git a/src/s_tir/schedule/analysis/reducer.cc b/src/s_tir/schedule/analysis/reducer.cc index ac15128f7c24..477b9b38fbdd 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 +#include #include #include "../utils.h" @@ -563,29 +564,21 @@ bool ReductionIterNotIndexOutputBuffer(const SBlock& block) { for (const MatchBufferRegion& region : block->match_buffers) { match_buffer_sources[region->buffer.get()] = region->source->buffer.get(); } - bool affected = false; - PreOrderVisit(block->body, [&](const ffi::ObjectRef& obj) { - if (affected) { - return false; - } - const auto* block_node = obj.as(); - if (block_node) { - for (const MatchBufferRegion& region : block_node->match_buffers) { - match_buffer_sources[region->buffer.get()] = region->source->buffer.get(); - } + auto visit_block = [&](const SBlock& nested_block) -> ffi::Expected { + for (const MatchBufferRegion& region : nested_block->match_buffers) { + match_buffer_sources[region->buffer.get()] = region->source->buffer.get(); } - // Inline AllocBufferNode statements (e.g. `T.local_scalar(...)` expansions) + return ffi::WalkResult::Advance(); + }; + auto visit_alloc = [&](const AllocBuffer& alloc) -> ffi::Expected { + // Inline AllocBuffer statements (e.g. `T.local_scalar(...)` expansions) // declare buffer-local scratch storage inside the block body; treat them // the same as block->alloc_buffers entries for the "write-without-signature" // check below. - if (const auto* alloc = obj.as()) { - buffer_allocated.insert(alloc->buffer.get()); - } - const auto* store = obj.as(); - if (!store) { - return true; - } - + buffer_allocated.insert(alloc->buffer.get()); + return ffi::WalkResult::Advance(); + }; + auto visit_store = [&](const BufferStore& store) -> ffi::Expected { bool write_is_covered_by_match_buffer = match_buffer_sources.count(store->buffer.get()) && buffer_written.count(match_buffer_sources.find(store->buffer.get())->second); @@ -593,17 +586,18 @@ bool ReductionIterNotIndexOutputBuffer(const SBlock& block) { buffer_allocated.count(store->buffer.get()), ValueError) << "The buffer \"" << store->buffer - << "\" is written in the block but is not in the block's signature nor is it covered by " - "a match_buffer"; + << "\" is written in the block but is not in the block's signature nor is it covered " + "by a match_buffer"; for (const PrimExpr& index : store->indices) { if (f_uses_reduction_block_var(index)) { - affected = true; - return false; + return ffi::WalkResult::Interrupt(ffi::VisitInterrupt(false)); } } - return false; - }); - return !affected; + return ffi::WalkResult::Skip(); + }; + auto result = ffi::StructuralWalk(block->body, visit_block, + visit_alloc, visit_store); + return result.has_value() ? result.value()->value.cast() : true; } class NoMatchedReducerError : public ScheduleError { diff --git a/src/s_tir/schedule/primitive/cache_index.cc b/src/s_tir/schedule/primitive/cache_index.cc index 4c04ac59409c..4afce512d1e0 100644 --- a/src/s_tir/schedule/primitive/cache_index.cc +++ b/src/s_tir/schedule/primitive/cache_index.cc @@ -18,6 +18,7 @@ */ #include #include +#include #include "../../../tirx/transform/replace_selected_expr.h" #include "../utils.h" @@ -172,21 +173,21 @@ class IndexInfoCollector : public StmtExprVisitor { // Record the final sub expr with repeat time greater than cse_thresh_ // In order to make the result stable, sort it by post order and then by complexity - PostOrderVisit(store->value, [&semantic_comp_done_by_stmt, this](const ffi::ObjectRef& node) { - if (auto prim = node.as()) { - PrimExpr this_expr = prim.value(); - for (auto& it : semantic_comp_done_by_stmt) { - if (it.second >= this->cse_thresh_ && EquivalentTerms(this_expr, it.first, true)) { - auto find_result = - std::find_if(this->exprs_.begin(), this->exprs_.end(), - [&](PrimExpr expr) { return expr.get() == it.first.get(); }); - if (find_result == this->exprs_.end()) { - this->exprs_.push_back(it.first); - } + auto walk_fn = [&semantic_comp_done_by_stmt, + this](const PrimExpr& this_expr) -> ffi::Expected { + for (auto& it : semantic_comp_done_by_stmt) { + if (it.second >= this->cse_thresh_ && EquivalentTerms(this_expr, it.first, true)) { + auto find_result = + std::find_if(this->exprs_.begin(), this->exprs_.end(), + [&](PrimExpr expr) { return expr.get() == it.first.get(); }); + if (find_result == this->exprs_.end()) { + this->exprs_.push_back(it.first); } } } - }); + return ffi::WalkResult::Advance(); + }; + ffi::StructuralWalk(store->value, walk_fn); auto cmp = [&](const PrimExpr& lhs, const PrimExpr& rhs) -> bool { return CalculateExprComplexity(lhs) > CalculateExprComplexity(rhs); }; @@ -235,9 +236,10 @@ ffi::Array MakeIndexCacheStage(IndexInfo* info, const ffi::String& stora // Collect the block vars in original index computation info->origin_block_vars.push_back({}); - PostOrderVisit(index_expr, [&info, &expr_index](const ffi::ObjectRef& node) { - if (auto var = node.as()) { - Var iter_var = var.value(); + auto collect_origin_var = [&info, + &expr_index](const Var& var) -> ffi::Expected { + if (auto prim_var = var.as()) { + Var iter_var = prim_var.value(); const ffi::Array& origin_block_var = info->origin_block_vars[expr_index]; auto find_result = std::find_if(origin_block_var.begin(), origin_block_var.end(), [&](Var it) { return it.get() == iter_var.get(); }); @@ -245,21 +247,25 @@ ffi::Array MakeIndexCacheStage(IndexInfo* info, const ffi::String& stora info->origin_block_vars[expr_index].push_back(iter_var); } } - }); + return ffi::WalkResult::Advance(); + }; + ffi::StructuralWalk(index_expr, collect_origin_var); // Collect the loop vars corresponding to collected block vars, // which will be used to create new loop vars std::vector iter_vars; - for (const Var& it : info->origin_block_vars[expr_index]) { - PostOrderVisit(info->var_binding.at(it), [/*&info,*/ &iter_vars](const ffi::ObjectRef& node) { - if (auto var = node.as()) { - Var iter_var = var.value(); - if (std::find_if(iter_vars.begin(), iter_vars.end(), - [&](Var it) { return it.get() == iter_var.get(); }) == iter_vars.end()) { - iter_vars.push_back(iter_var); - } + auto collect_iter_var = [&iter_vars](const Var& var) -> ffi::Expected { + if (auto prim_var = var.as()) { + Var iter_var = prim_var.value(); + if (std::find_if(iter_vars.begin(), iter_vars.end(), + [&](Var it) { return it.get() == iter_var.get(); }) == iter_vars.end()) { + iter_vars.push_back(iter_var); } - }); + } + return ffi::WalkResult::Advance(); + }; + for (const Var& it : info->origin_block_vars[expr_index]) { + ffi::StructuralWalk(info->var_binding.at(it), collect_iter_var); } PrimType data_ty = index_expr.ty(); diff --git a/src/s_tir/schedule/primitive/cache_read_write.cc b/src/s_tir/schedule/primitive/cache_read_write.cc index d67d2486d58c..878941a4bb6e 100644 --- a/src/s_tir/schedule/primitive/cache_read_write.cc +++ b/src/s_tir/schedule/primitive/cache_read_write.cc @@ -18,6 +18,7 @@ */ #include +#include #include @@ -2332,13 +2333,12 @@ StmtSRef ReIndex(ScheduleState self, const StmtSRef& block_sref, int buffer_inde // Collect block iters appearing in the original_indices std::unordered_set covered; + auto walk_fn = [&covered](const Var& var) -> ffi::Expected { + covered.insert(var); + return ffi::WalkResult::Advance(); + }; for (const PrimExpr& index : original_indices) { - PreOrderVisit(index, [&](const ffi::ObjectRef& obj) -> bool { - if (auto var = obj.as()) { - covered.insert(var.value()); - } - return true; - }); + ffi::StructuralWalk(index, walk_fn); } // Step 2. Creating CacheStageInfo diff --git a/src/s_tir/schedule/primitive/for_kind.cc b/src/s_tir/schedule/primitive/for_kind.cc index c2c7c1c887f3..833877af216d 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 +#include #include "../utils.h" @@ -124,18 +125,16 @@ void CheckLoopParallelizableInBlock(const ScheduleState& self, ForKind for_kind, */ void CheckParallelizability(const ScheduleState& self, const For& loop, ForKind for_kind, runtime::ThreadScope thread_scope) { - PreOrderVisit(loop, [&](const ffi::ObjectRef& node) { - if (const auto* realize = node.as()) { - // If this block doesn't have corresponding StmtSRef in the schedule state, it must be a block - // inside `tirx.init()`. We don't check the condition for such blocks. - if (!self->stmt2ref.count(realize->block.get())) { - return false; - } - CheckLoopParallelizableInBlock(self, for_kind, loop->loop_var, - ffi::GetRef(realize), thread_scope); + auto walk_fn = [&](const SBlockRealize& realize) -> ffi::Expected { + // If this block doesn't have corresponding StmtSRef in the schedule state, it must be a + // block inside `tirx.init()`. We don't check the condition for such blocks. + if (!self->stmt2ref.count(realize->block.get())) { + return ffi::WalkResult::Skip(); } - return true; - }); + CheckLoopParallelizableInBlock(self, for_kind, loop->loop_var, realize, thread_scope); + return ffi::WalkResult::Advance(); + }; + ffi::StructuralWalk(loop, walk_fn); } /*! diff --git a/src/s_tir/schedule/primitive/layout_transformation.cc b/src/s_tir/schedule/primitive/layout_transformation.cc index 696821f3227e..4fc6ea3dc835 100644 --- a/src/s_tir/schedule/primitive/layout_transformation.cc +++ b/src/s_tir/schedule/primitive/layout_transformation.cc @@ -19,6 +19,7 @@ #include #include +#include #include #include @@ -1273,21 +1274,22 @@ IterVarType DetectNewBlockIterType( const std::unordered_map& block_iter_type_map) { IterVarType result{kOpaque}; bool found = false; - PostOrderVisit(expr, [&](const ffi::ObjectRef& obj) { - if (auto var = obj.as()) { - auto it = block_iter_type_map.find(var.value().get()); + auto walk_fn = [&](const Var& var) -> ffi::Expected { + if (auto prim_var = var.as()) { + auto it = block_iter_type_map.find(prim_var.value().get()); if (it != block_iter_type_map.end()) { if (!found) { found = true; result = it->second; } else if (result != it->second) { result = kOpaque; - return false; + return ffi::WalkResult::Interrupt(); } } } - return true; - }); + return ffi::WalkResult::Advance(); + }; + ffi::StructuralWalk(expr, walk_fn); return result; } diff --git a/src/s_tir/schedule/primitive/pad_einsum.cc b/src/s_tir/schedule/primitive/pad_einsum.cc index e07d1c550c02..6f08e80fb4f4 100644 --- a/src/s_tir/schedule/primitive/pad_einsum.cc +++ b/src/s_tir/schedule/primitive/pad_einsum.cc @@ -18,6 +18,7 @@ */ #include +#include #include #include "../utils.h" @@ -423,13 +424,14 @@ void PadEinsum(ScheduleState self, const StmtSRef& block_sref, const ffi::Array< // Step 4. Find out the block of our interest int pos = -1; for (int i = 0; i < static_cast(scope_body.size()); ++i) { - bool found = false; - PostOrderVisit(scope_body[i], [&found, &block](const ffi::ObjectRef& node) { + auto walk_fn = [&block](const SBlock& node) -> ffi::Expected { if (node.get() == block) { - found = true; + return ffi::WalkResult::Interrupt(ffi::VisitInterrupt(true)); } - }); - if (found) { + return ffi::WalkResult::Advance(); + }; + auto result = ffi::StructuralWalk(scope_body[i], walk_fn); + if (result.has_value() && result.value()->value.cast()) { pos = i; break; } diff --git a/src/s_tir/schedule/primitive/read_write_at.cc b/src/s_tir/schedule/primitive/read_write_at.cc index b7352c272afd..ef40ba230f42 100644 --- a/src/s_tir/schedule/primitive/read_write_at.cc +++ b/src/s_tir/schedule/primitive/read_write_at.cc @@ -18,6 +18,7 @@ */ #include +#include #include #include @@ -184,11 +185,7 @@ struct ReadWriteAtImpl { bool r_visited = false; bool w_visited = false; auto f_visit = [this, &relaxed_regions, &r_visited, &w_visited, - &scope](const ffi::ObjectRef& obj) -> bool { - const SBlockRealizeNode* realize = obj.as(); - if (realize == nullptr) { - return true; - } + &scope](const SBlockRealize& realize) -> ffi::Expected { const SBlockNode* block = realize->block.get(); bool has_r = HasBuffer(block->reads, src_); bool has_w = HasBuffer(block->writes, src_); @@ -203,12 +200,12 @@ struct ReadWriteAtImpl { /*low_inclusive=*/ffi::GetRef(self_->stmt2ref.at(block)->parent), /*high_exclusive=*/loop_sref_, /*extra_relax_scope=*/scope)), - /*bindings=*/GetBindings(ffi::GetRef(realize)), + /*bindings=*/GetBindings(realize), /*relaxed_regions=*/&relaxed_regions); } - return false; + return ffi::WalkResult::Skip(); }; - PreOrderVisit(subtrees[i], f_visit); + ffi::StructuralWalk(subtrees[i], f_visit); if (r_visited) { r_pos.push_back(i); } diff --git a/src/s_tir/schedule/trace.cc b/src/s_tir/schedule/trace.cc index a8272f0c9412..acaeba7aca43 100644 --- a/src/s_tir/schedule/trace.cc +++ b/src/s_tir/schedule/trace.cc @@ -17,6 +17,7 @@ * under the License. */ #include +#include #include #include @@ -583,11 +584,11 @@ Trace TraceNode::Simplified(bool remove_postproc) const { used_rvs.insert(obj.as()); continue; } else if (auto prim_expr = obj.as()) { - PostOrderVisit(*prim_expr, [&used_rvs](const ffi::ObjectRef& obj) -> void { - if (obj.as()) { - used_rvs.insert(obj.get()); - } - }); + auto walk_fn = [&used_rvs](const Var& var) -> ffi::Expected { + used_rvs.insert(var.get()); + return ffi::WalkResult::Advance(); + }; + ffi::StructuralWalk(*prim_expr, walk_fn); } } } diff --git a/src/s_tir/transform/lower_cross_thread_reduction.cc b/src/s_tir/transform/lower_cross_thread_reduction.cc index 474b23b2f929..47c87dfeb9ba 100644 --- a/src/s_tir/transform/lower_cross_thread_reduction.cc +++ b/src/s_tir/transform/lower_cross_thread_reduction.cc @@ -22,6 +22,7 @@ */ #include #include +#include #include #include #include @@ -80,15 +81,13 @@ bool IsBoundToThreadIdx(const ForNode* loop) { bool IsDominantBlock(const SBlock& scope_block, const SBlock& block) { // Step 1. Count the number of writers for each buffer written by the scope block. std::unordered_map buffer_writer_cnt; - PreOrderVisit(scope_block->body, [&buffer_writer_cnt](const ffi::ObjectRef& obj) { - if (const auto* block = obj.as()) { - for (const BufferRegion& buffer_region : block->writes) { - ++buffer_writer_cnt[buffer_region->buffer.get()]; - } - return false; + auto walk_fn = [&buffer_writer_cnt](const SBlock& block) -> ffi::Expected { + for (const BufferRegion& buffer_region : block->writes) { + ++buffer_writer_cnt[buffer_region->buffer.get()]; } - return true; - }); + return ffi::WalkResult::Skip(); + }; + ffi::StructuralWalk(scope_block->body, walk_fn); // Step 2. Check whether `block` is the only writer of its outputs. for (const BufferRegion& buffer_region : block->writes) { TVM_FFI_ICHECK(buffer_writer_cnt.count(buffer_region->buffer.get())); @@ -479,31 +478,33 @@ Stmt TransformReductionBlock(const SBlockRealizeNode* realize, for (const ForNode* reduction_loop : reduction_loops) { reduction_loop_vars.insert(reduction_loop->loop_var.get()); } - PostOrderVisit(realize->predicate, - [&wb_predicate, &reduction_loop_vars](const ffi::ObjectRef& obj) { - if (const auto* and_node = obj.as()) { - ffi::Array sub_exprs = {and_node->a, and_node->b}; - for (PrimExpr sub_expr : sub_exprs) { - if (sub_expr->IsInstance()) { - continue; - } - bool is_reduction = [sub_expr, &reduction_loop_vars]() { - ffi::Array vars = UndefinedVars(sub_expr); - for (Var var : vars) { - if (reduction_loop_vars.find(var.get()) != reduction_loop_vars.end()) { - return true; - } - } - return false; - }(); - if (!is_reduction) { - wb_predicate = wb_predicate && sub_expr; - } - } - return true; - } - return false; - }); + std::unordered_set visited_predicate_nodes; + auto walk_fn = [&wb_predicate, &reduction_loop_vars, &visited_predicate_nodes]( + const And& and_expr) -> ffi::Expected { + if (!visited_predicate_nodes.insert(and_expr.get()).second) { + return ffi::WalkResult::Advance(); + } + ffi::Array sub_exprs = {and_expr->a, and_expr->b}; + for (PrimExpr sub_expr : sub_exprs) { + if (sub_expr->IsInstance()) { + continue; + } + bool is_reduction = [sub_expr, &reduction_loop_vars]() { + ffi::Array vars = UndefinedVars(sub_expr); + for (Var var : vars) { + if (reduction_loop_vars.find(var.get()) != reduction_loop_vars.end()) { + return true; + } + } + return false; + }(); + if (!is_reduction) { + wb_predicate = wb_predicate && sub_expr; + } + } + return ffi::WalkResult::Advance(); + }; + ffi::StructuralWalk(realize->predicate, walk_fn); if (wb_buffers[0].scope() != "local") { for (const ForNode* loop : reduction_loops) { if (loop->thread_binding.has_value()) { @@ -711,18 +712,16 @@ class CrossThreadReductionTransformer : public StmtMutator { // Condition 5. The block should be the last block under the first reduction-related loop. bool visit = false; - PreOrderVisit(ffi::GetRef(reduction_loops[0]), [block, &visit](const ffi::ObjectRef& obj) { - if (const auto* realize = obj.as()) { - TVM_FFI_CHECK(!visit, ValueError) - << "Cross-thread reduction cannot be applied when the reduction " - "block isn't the last block under its first reduction-related loop"; - if (realize->block.get() == block) { - visit = true; - } - return false; + auto walk_fn = [block, &visit](const SBlockRealize& realize) -> ffi::Expected { + TVM_FFI_CHECK(!visit, ValueError) + << "Cross-thread reduction cannot be applied when the reduction " + "block isn't the last block under its first reduction-related loop"; + if (realize->block.get() == block) { + visit = true; } - return true; - }); + return ffi::WalkResult::Skip(); + }; + ffi::StructuralWalk(ffi::GetRef(reduction_loops[0]), walk_fn); return std::make_tuple(n_bound_reduction_loops, // std::move(reducer), // std::move(reduction_buffers), // diff --git a/src/s_tir/transform/memhammer_intermediate_stage.cc b/src/s_tir/transform/memhammer_intermediate_stage.cc index 20272f87d2c6..3ce33fb6f8de 100644 --- a/src/s_tir/transform/memhammer_intermediate_stage.cc +++ b/src/s_tir/transform/memhammer_intermediate_stage.cc @@ -17,6 +17,7 @@ * under the License. */ #include +#include #include "memhammer_rewrite_rule.h" @@ -282,26 +283,25 @@ std::pair InsertCacheStage(Stmt stmt, bool is_write_cache, ffi::S arith::Analyzer analyzer; const TensorLoadNode* target_buffer_load = nullptr; if (is_write_cache) { - tirx::PreOrderVisit(stmt, [&](const ffi::ObjectRef& obj) { - if (const auto* buffer_load = obj.as()) { - if (buffer_load->source.as_or_throw().scope() == "wmma.accumulator" || - buffer_load->source.as_or_throw().scope() == "m16n8k8.matrixC") { - if (target_buffer_load == nullptr) { - target_buffer_load = buffer_load; - } else { - TVM_FFI_ICHECK(target_buffer_load->source.as_or_throw().same_as( - buffer_load->source.as_or_throw())) - << "More than one target buffer found"; - TVM_FFI_ICHECK(target_buffer_load->indices.size() == buffer_load->indices.size()); - for (size_t i = 0; i < target_buffer_load->indices.size(); i++) { - TVM_FFI_ICHECK( - analyzer->CanProveEqual(target_buffer_load->indices[i], buffer_load->indices[i])); - } + auto walk_fn = [&](const TensorLoad& buffer_load) -> ffi::Expected { + if (buffer_load->source.as_or_throw().scope() == "wmma.accumulator" || + buffer_load->source.as_or_throw().scope() == "m16n8k8.matrixC") { + if (target_buffer_load == nullptr) { + target_buffer_load = buffer_load.get(); + } else { + TVM_FFI_ICHECK(target_buffer_load->source.as_or_throw().same_as( + buffer_load->source.as_or_throw())) + << "More than one target buffer found"; + TVM_FFI_ICHECK(target_buffer_load->indices.size() == buffer_load->indices.size()); + for (size_t i = 0; i < target_buffer_load->indices.size(); i++) { + TVM_FFI_ICHECK( + analyzer->CanProveEqual(target_buffer_load->indices[i], buffer_load->indices[i])); } } } - return true; - }); + return ffi::WalkResult::Advance(); + }; + ffi::StructuralWalk(stmt, walk_fn); TVM_FFI_ICHECK(target_buffer_load); } diff --git a/src/s_tir/transform/memhammer_tensorcore_rewrite.cc b/src/s_tir/transform/memhammer_tensorcore_rewrite.cc index 2de9369125d9..29092d821b6e 100644 --- a/src/s_tir/transform/memhammer_tensorcore_rewrite.cc +++ b/src/s_tir/transform/memhammer_tensorcore_rewrite.cc @@ -18,6 +18,7 @@ */ #include +#include #include #include "./memhammer_rewrite_rule.h" @@ -228,17 +229,17 @@ Stmt RewriteWmmaStore(Stmt stmt) { // TODO(tian): the assumption that the RHS of BufferStore is TensorLoad may not be accurate const BufferStoreNode* buf_store = TVM_TYPE_AS(body, BufferStoreNode); const TensorLoadNode* buf_load = nullptr; - PostOrderVisit(buf_store->value, [&](const ffi::ObjectRef& obj) { - const TensorLoadNode* load = obj.as(); - if (load && load->source.as_or_throw().scope() == "wmma.accumulator") { + auto walk_fn = [&](const TensorLoad& load) -> ffi::Expected { + if (load->source.as_or_throw().scope() == "wmma.accumulator") { TVM_FFI_ICHECK(buf_load == nullptr || buf_load->source.as_or_throw().same_as( load->source.as_or_throw())) << "More than one source buffer of wmma accumulator found"; - buf_load = load; + buf_load = load.get(); } - return true; - }); + return ffi::WalkResult::Advance(); + }; + ffi::StructuralWalk(buf_store->value, walk_fn); BufferVar src_buffer = buf_load->source.as_or_throw(); BufferVar tgt_buffer = buf_store->buffer; @@ -439,17 +440,17 @@ Stmt RewriteMmaStore(Stmt stmt) { // Step 2. Find matrixC buffer const BufferStoreNode* buf_store = TVM_TYPE_AS(body, BufferStoreNode); const TensorLoadNode* buf_load = nullptr; - PostOrderVisit(buf_store->value, [&](const ffi::ObjectRef& obj) { - const TensorLoadNode* load = obj.as(); - if (load && load->source.as_or_throw().scope() == "m16n8k8.matrixC") { + auto walk_fn = [&](const TensorLoad& load) -> ffi::Expected { + if (load->source.as_or_throw().scope() == "m16n8k8.matrixC") { TVM_FFI_ICHECK(buf_load == nullptr || buf_load->source.as_or_throw().same_as( load->source.as_or_throw())) << "More than one source buffer of mma accumulator found"; - buf_load = load; + buf_load = load.get(); } - return true; - }); + return ffi::WalkResult::Advance(); + }; + ffi::StructuralWalk(buf_store->value, walk_fn); // Step 3. Create new mma body // We have the assumption that two innermost loops are the 8 * 8 loop generated by diff --git a/src/s_tir/transform/renew_defs.cc b/src/s_tir/transform/renew_defs.cc index 212fd3c2244d..86fe153054ce 100644 --- a/src/s_tir/transform/renew_defs.cc +++ b/src/s_tir/transform/renew_defs.cc @@ -23,6 +23,7 @@ */ #include +#include #include #include #include @@ -59,14 +60,14 @@ class RenewDefMutator : public StmtExprMutator { for (const auto& param : func->params) { if (auto opt_buffer = param.as()) { const BufferVar& buffer = opt_buffer.value(); + auto walk_fn = [&generator](const Var& var) -> ffi::Expected { + if (generator.remap_.count(var) == 0) { + generator.ReDefineVar(var); + } + return ffi::WalkResult::Advance(); + }; for (const PrimExpr& e : buffer->shape) { - PostOrderVisit(e, [&generator](const ffi::ObjectRef& obj) { - if (auto var = obj.as()) { - if (generator.remap_.count(var.value()) == 0) { - generator.ReDefineVar(var.value()); - } - } - }); + ffi::StructuralWalk(e, walk_fn); } } } diff --git a/src/te/operation/compute_op.cc b/src/te/operation/compute_op.cc index 4bd50d15bd93..1600541d8147 100644 --- a/src/te/operation/compute_op.cc +++ b/src/te/operation/compute_op.cc @@ -23,6 +23,7 @@ */ #include +#include #include #include #include @@ -163,16 +164,18 @@ TVM_FFI_STATIC_INIT_BLOCK() { ffi::Array ComputeOpNode::InputTensors() const { ffi::Array ret; std::unordered_set visited; - auto visit = [&ret, &visited](const PrimExpr& e) { - tirx::PostOrderVisit(e, [&ret, &visited](const ffi::ObjectRef& n) { - if (auto call = n.as(); call.has_value() && IsTensorLoad(call.value())) { - Tensor t = GetTensorFromLoad(call.value()); - if (!visited.count(t)) { - ret.push_back(t); - visited.insert(t); - } + auto walk_fn = [&ret, &visited](const Call& call) -> ffi::Expected { + if (IsTensorLoad(call)) { + Tensor t = GetTensorFromLoad(call); + if (!visited.count(t)) { + ret.push_back(t); + visited.insert(t); } - }); + } + return ffi::WalkResult::Advance(); + }; + auto visit = [&walk_fn](const PrimExpr& e) { + ffi::StructuralWalk(e, walk_fn); }; for (const PrimExpr& e : body) { if (const auto* reduce = e.as()) { diff --git a/src/te/tensor.cc b/src/te/tensor.cc index a617d54bd4ff..15d644a9fec8 100644 --- a/src/te/tensor.cc +++ b/src/te/tensor.cc @@ -21,6 +21,8 @@ * \file tensor.cc */ #include +#include +#include #include #include #include @@ -29,6 +31,22 @@ namespace tvm { namespace te { +namespace { + +TVMFFIAny TensorVisit(ffi::StructuralVisitorObj*, ffi::AnyView) noexcept { + return ffi::AnyView(nullptr).CopyToTVMFFIAny(); +} + +TVMFFIAny TensorMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept { + return ffi::Unchanged().CopyToTVMFFIAny(); +} + +TVMFFIAny TensorMaybeInplaceMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept { + return ffi::Unchanged().CopyToTVMFFIAny(); +} + +} // namespace + void TensorNode::RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef() @@ -38,7 +56,15 @@ void TensorNode::RegisterReflection() { .def_ro("value_index", &TensorNode::value_index); } -TVM_FFI_STATIC_INIT_BLOCK() { TensorNode::RegisterReflection(); } +TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; + TensorNode::RegisterReflection(); + refl::TypeAttrDef() + .attr(refl::type_attr::kStructuralVisit, reinterpret_cast(&TensorVisit)) + .attr(refl::type_attr::kStructuralMutate, reinterpret_cast(&TensorMutate)) + .attr(refl::type_attr::kStructuralMaybeInplaceMutate, + reinterpret_cast(&TensorMaybeInplaceMutate)); +} IterVar thread_axis(Range dom, std::string tag) { return IterVar(dom, PrimVar(tag, dom.defined() ? dom->extent.ty() : PrimType::Int(32)), diff --git a/src/tirx/ir/index_map.cc b/src/tirx/ir/index_map.cc index e4e473990d64..1f53c6395c29 100644 --- a/src/tirx/ir/index_map.cc +++ b/src/tirx/ir/index_map.cc @@ -25,6 +25,7 @@ #include #include #include +#include #include #include #include @@ -362,23 +363,25 @@ IndexMap IndexMap::RenameVariables( if (f_name_map != nullptr) { // Collect variables with pre-defined names provided by f_name_map. std::unordered_set visited; + auto walk_fn = [&](const Var& var) -> ffi::Expected { + auto prim_var = var.as(); + if (!prim_var) { + return ffi::WalkResult::Advance(); + } + if (!visited.insert(prim_var.value().get()).second) { + return ffi::WalkResult::Advance(); + } + if (ffi::Optional opt_name = f_name_map(prim_var.value()); + opt_name.has_value()) { + ffi::String name = opt_name.value(); + TVM_FFI_ICHECK(!name_supply->ContainsName(name, /*add_prefix=*/false)); + name_supply->ReserveName(name, /*add_prefix=*/false); + var_remap.Set(prim_var.value(), PrimVar(name, prim_var.value().ty())); + } + return ffi::WalkResult::Advance(); + }; std::for_each(n->final_indices.begin(), n->final_indices.end(), [&](const PrimExpr& expr) { - PostOrderVisit(expr, [&](const ffi::ObjectRef& obj) { - auto var = obj.as(); - if (!var) { - return; - } - if (visited.count(obj.get())) { - return; - } - visited.emplace(obj.get()); - if (ffi::Optional opt_name = f_name_map(var.value()); opt_name.has_value()) { - ffi::String name = opt_name.value(); - TVM_FFI_ICHECK(!name_supply->ContainsName(name, /*add_prefix=*/false)); - name_supply->ReserveName(name, /*add_prefix=*/false); - var_remap.Set(var.value(), PrimVar(name, var.value().ty())); - } - }); + ffi::StructuralWalk(expr, walk_fn); }); } diff --git a/src/tirx/ir/stmt_functor.cc b/src/tirx/ir/stmt_functor.cc index 251f9d217ea7..45cd939e915b 100644 --- a/src/tirx/ir/stmt_functor.cc +++ b/src/tirx/ir/stmt_functor.cc @@ -743,43 +743,7 @@ Stmt StmtMutator::VisitStmt_(const tirx::TilePrimitiveCallNode* op) { } } -// Implementations of PostOrderVisit and Substitute -class IRApplyVisit : public StmtExprVisitor { - public: - explicit IRApplyVisit(std::function f) : f_(f) {} - - void VisitExpr(const Expr& node) final { - if (visited_.count(node.get()) != 0) return; - visited_.insert(node.get()); - ExprVisitor::VisitExpr(node); - f_(node); - } - - void VisitStmt(const Stmt& node) final { - if (visited_.count(node.get()) != 0) return; - visited_.insert(node.get()); - StmtVisitor::VisitStmt(node); - f_(node); - } - - void VisitBufferDef(const BufferVar& buffer, bool alloc_data) override {} - void VisitBufferUse(const BufferVar& buffer) override {} - - private: - std::function f_; - std::unordered_set visited_; -}; - -void PostOrderVisit(const ffi::ObjectRef& node, std::function fvisit) { - if (node.as()) { - IRApplyVisit visitor(fvisit); - visitor(node.as_or_throw()); - } else { - IRApplyVisit visitor(fvisit); - visitor(node.as_or_throw()); - } -} - +// Implementations of Substitute class IRSubstitute : public StmtExprMutator { public: explicit IRSubstitute(std::function(const Var&)> vmap) : vmap_(vmap) {} @@ -854,48 +818,6 @@ Expr Substitute(Expr expr, std::function(const Var&)> vmap) return IRSubstitute(std::move(vmap))(std::move(expr)); } -void PreOrderVisit(const ffi::ObjectRef& stmt_or_expr, - const std::function& fvisit) { - class PreOrderVisitor : public StmtExprVisitor { - public: - explicit PreOrderVisitor(const std::function& f) : f_(f) {} - - private: - void VisitExpr(const Expr& expr) final { - const ExprNode* p_expr = expr.get(); - if (visited_.count(p_expr) == 0) { - visited_.insert(p_expr); - if (f_(expr)) { - ExprVisitor::VisitExpr(expr); - } - } - } - - void VisitStmt(const Stmt& stmt) final { - const StmtNode* p_stmt = stmt.get(); - if (visited_.count(p_stmt) == 0) { - visited_.insert(p_stmt); - if (f_(stmt)) { - StmtVisitor::VisitStmt(stmt); - } - } - } - - const std::function& f_; - std::unordered_set visited_; - }; - - PreOrderVisitor visitor(fvisit); - if (auto stmt = stmt_or_expr.as()) { - visitor(stmt.value()); - } else if (auto expr = stmt_or_expr.as()) { - visitor(expr.value()); - } else { - TVM_FFI_THROW(InternalError) << "PreOrderVisit does not accept object with type: " - << stmt_or_expr->GetTypeKey(); - } -} - class IRSubstituteWithDataTypeLegalization : public DataTypeLegalizer { public: explicit IRSubstituteWithDataTypeLegalization(std::function(const Var&)> vmap) @@ -951,22 +873,14 @@ PrimExpr SubstituteWithDataTypeLegalization( TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef() - .def("tirx.PostOrderVisit", - [](ffi::ObjectRef node, ffi::Function f) { - tirx::PostOrderVisit(node, [f](const ffi::ObjectRef& n) { f(n); }); - }) - .def("tirx.PreOrderVisit", - [](ffi::ObjectRef node, ffi::Function f) { - tirx::PreOrderVisit(node, [f](const ffi::ObjectRef& n) { return f(n).cast(); }); - }) - .def("tirx.Substitute", [](ffi::ObjectRef node, ffi::Map vmap) -> ffi::ObjectRef { - if (node->IsInstance()) { - return Substitute(node.as_or_throw(), vmap); - } else { - return Substitute(node.as_or_throw(), vmap); - } - }); + refl::GlobalDef().def("tirx.Substitute", + [](ffi::ObjectRef node, ffi::Map vmap) -> ffi::ObjectRef { + if (node->IsInstance()) { + return Substitute(node.as_or_throw(), vmap); + } else { + return Substitute(node.as_or_throw(), vmap); + } + }); } } // namespace tirx diff --git a/src/tirx/ir/tir_visitor_with_path.h b/src/tirx/ir/tir_visitor_with_path.h index b14ad713f32f..13007a32d2aa 100644 --- a/src/tirx/ir/tir_visitor_with_path.h +++ b/src/tirx/ir/tir_visitor_with_path.h @@ -24,6 +24,7 @@ #ifndef TVM_TIRX_IR_TIR_VISITOR_WITH_PATH_H_ #define TVM_TIRX_IR_TIR_VISITOR_WITH_PATH_H_ +#include #include #include #include @@ -273,13 +274,13 @@ class TIRVisitorWithPath : protected ExprFunctorAttr("shape"); for (size_t i = 0; i < buf->shape.size(); i++) { auto dim_path = shape_path->ArrayItem(i); - PostOrderVisit(buf->shape[i], [this, &context, &dim_path](const ffi::ObjectRef& obj) { - if (auto opt = obj.as()) { - if (auto var_def = WithDefIfUndefined(opt.value(), dim_path)) { - context.push_back(std::move(var_def).value()); - } + auto walk_fn = [this, &context, &dim_path](const Var& var) -> ffi::Expected { + if (auto var_def = WithDefIfUndefined(var, dim_path)) { + context.push_back(std::move(var_def).value()); } - }); + return ffi::WalkResult::Advance(); + }; + ffi::StructuralWalk(buf->shape[i], walk_fn); } auto strides_path = path->Attr("strides"); diff --git a/src/tirx/script/printer/buffer.cc b/src/tirx/script/printer/buffer.cc index fa5b0b605832..c3c456f390fd 100644 --- a/src/tirx/script/printer/buffer.cc +++ b/src/tirx/script/printer/buffer.cc @@ -40,22 +40,25 @@ ffi::Map BufferAttrs( // Step 0. Set up statistics std::unordered_map use_count; - auto update_use_count = [&](const Expr& e) { - tirx::PostOrderVisit(e, [&](const ffi::ObjectRef& n) { - if (const VarNode* var = n.as()) { - ++use_count[var]; + std::unordered_set def_seen; + auto count_buffer_var = [&](const Var& var, + TVMFFIDefRegionKind kind) -> ffi::Expected { + if (kind != kTVMFFIDefRegionKindNone) { + if (!def_seen.insert(var.get()).second) { + return ffi::WalkResult::Skip(); } - }); + return ffi::WalkResult::Advance(); + } + ++use_count[var.get()]; + return ffi::WalkResult::Advance(); }; - update_use_count(buffer->elem_offset); + ffi::StructuralWalk(buffer, count_buffer_var); if (data.has_value()) { - update_use_count(data.value()); - } - for (const PrimExpr& e : buffer->strides) { - update_use_count(e); - } - for (const PrimExpr& e : buffer->shape) { - update_use_count(e); + auto count_data_var = [&](const Var& var) -> ffi::Expected { + ++use_count[var.get()]; + return ffi::WalkResult::Advance(); + }; + ffi::StructuralWalk(data.value(), count_data_var); } auto is_new_var = [&](const Expr& e) { return e->IsInstance() && !d->IsVarDefined(e); }; auto add_out_of_line_var_def = [&](const Var& var, const AccessPath& var_p) { @@ -90,16 +93,15 @@ ffi::Map BufferAttrs( bool contains_new_var = false; bool contains_compound_shape_var = false; std::unordered_set vars_in_shape; - tirx::PostOrderVisit(e, [&](const ffi::ObjectRef& obj) { - if (const auto* var_node = obj.as()) { - Var var = ffi::GetRef(var_node); - vars_in_shape.insert(var); - contains_new_var = - contains_new_var || !d->IsVarDefined(var) || stringify_shape_vars.count(var); - contains_compound_shape_var = - contains_compound_shape_var || stringify_compound_shape_vars.count(var); - } - }); + auto walk_fn = [&](const Var& var) -> ffi::Expected { + vars_in_shape.insert(var); + contains_new_var = + contains_new_var || !d->IsVarDefined(var) || stringify_shape_vars.count(var); + contains_compound_shape_var = + contains_compound_shape_var || stringify_compound_shape_vars.count(var); + return ffi::WalkResult::Advance(); + }; + ffi::StructuralWalk(e, walk_fn); if (is_new_var(e)) { add_out_of_line_var_def(e.as_or_throw(), e_p); } diff --git a/src/tirx/script/printer/function.cc b/src/tirx/script/printer/function.cc index 554dba11c549..a461cddfc092 100644 --- a/src/tirx/script/printer/function.cc +++ b/src/tirx/script/printer/function.cc @@ -93,21 +93,20 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) std::unordered_set stringify_shape_vars; std::unordered_set stringify_compound_shape_vars; std::unordered_set shape_vars; + auto walk_fn = [&](const tirx::Var& shape_var) -> ffi::Expected { + shape_vars.insert(shape_var); + bool is_type_var = type_vars.count(shape_var.get()); + if (!use_postponed_annotations && !bound_signature_vars.count(shape_var) && + !is_type_var) { + stringify_shape_vars.insert(shape_var); + } + if (!use_postponed_annotations && is_type_var) { + stringify_compound_shape_vars.insert(shape_var); + } + return ffi::WalkResult::Advance(); + }; for (const PrimExpr& shape : buffer->shape) { - tirx::PostOrderVisit(shape, [&](const ffi::ObjectRef& obj) { - if (const auto* shape_var_node = obj.as()) { - tirx::Var shape_var = ffi::GetRef(shape_var_node); - shape_vars.insert(shape_var); - bool is_type_var = type_vars.count(shape_var.get()); - if (!use_postponed_annotations && !bound_signature_vars.count(shape_var) && - !is_type_var) { - stringify_shape_vars.insert(shape_var); - } - if (!use_postponed_annotations && is_type_var) { - stringify_compound_shape_vars.insert(shape_var); - } - } - }); + ffi::StructuralWalk(shape, walk_fn); } IdDoc lhs = DefineBuffer(buffer, *f, d); ExprDoc annotation = diff --git a/src/tirx/script/printer/utils.h b/src/tirx/script/printer/utils.h index 4807a37e8451..83ddf213c60e 100644 --- a/src/tirx/script/printer/utils.h +++ b/src/tirx/script/printer/utils.h @@ -20,6 +20,7 @@ #define TVM_SCRIPT_PRINTER_TIR_UTILS_H_ #include +#include #include #include #include @@ -117,21 +118,22 @@ inline void AsDocBody(const tirx::Stmt& stmt, AccessPath p, TIRFrameNode* f, con if (const auto* seq_stmt = stmt.as()) { ffi::Array body = seq_stmt->seq; auto value_refs_buffer = [](const PrimExpr& value, const tirx::BufferVar& buffer) { - bool found = false; - tirx::PostOrderVisit(value, [&](const ffi::ObjectRef& node) { - if (const auto* load = node.as()) { - if (load->source.as_or_throw().same_as(buffer)) { - found = true; - } - } else if (const auto* call = node.as()) { - if (call->op.same_as(tirx::builtin::masked_load()) && !call->args.empty()) { - if (auto var = call->args[0].as(); var && var.value().same_as(buffer.var())) { - found = true; - } + auto visit_load = [&](const TensorLoad& load) -> ffi::Expected { + if (load->source.as_or_throw().same_as(buffer)) { + return ffi::WalkResult::Interrupt(ffi::VisitInterrupt(true)); + } + return ffi::WalkResult::Advance(); + }; + auto visit_call = [&](const Call& call) -> ffi::Expected { + if (call->op.same_as(tirx::builtin::masked_load()) && !call->args.empty()) { + if (auto var = call->args[0].as(); var && var.value().same_as(buffer.var())) { + return ffi::WalkResult::Interrupt(ffi::VisitInterrupt(true)); } } - }); - return found; + return ffi::WalkResult::Advance(); + }; + auto result = ffi::StructuralWalk(value, visit_load, visit_call); + return result.has_value() ? result.value()->value.cast() : false; }; for (int i = 0, n = body.size(); i < n;) { diff --git a/src/tirx/transform/ir_utils.cc b/src/tirx/transform/ir_utils.cc index e41223d819fc..6e7e44b780f0 100644 --- a/src/tirx/transform/ir_utils.cc +++ b/src/tirx/transform/ir_utils.cc @@ -26,6 +26,7 @@ #include #include #include +#include #include #include #include @@ -129,10 +130,12 @@ class IRConvertSSA final : public StmtExprMutator { defined_.insert(var_ptr); } }; + auto walk_fn = [&](const Var& var) -> ffi::Expected { + check_var(var); + return ffi::WalkResult::Advance(); + }; for (const auto& dim : buffer.value()->shape) { - PostOrderVisit(dim, [&](const ffi::ObjectRef& obj) { - if (auto var = obj.as()) check_var(var.value()); - }); + ffi::StructuralWalk(dim, walk_fn); } for (const auto& stride : buffer.value()->strides) { if (auto var = stride.as()) check_var(var.value()); @@ -793,9 +796,10 @@ ffi::Optional ConditionalBoundsContext::TrySolveCondition e->IsInstance() || e->IsInstance()) { bool is_simple = true; std::vector cand_vars; - PostOrderVisit(e, [&cand_vars, &is_simple, &e](const ffi::ObjectRef& obj) { + auto walk_fn = [&cand_vars, &is_simple, + &e](const PrimExpr& obj) -> ffi::Expected { if (obj.same_as(e)) { - return; + return ffi::WalkResult::Advance(); } else if (const VarNode* var = obj.as()) { PrimType var_ty = var->ty.as_or_throw(); if (var_ty.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt)) { @@ -806,7 +810,9 @@ ffi::Optional ConditionalBoundsContext::TrySolveCondition obj->IsInstance() || obj->IsInstance() || obj->IsInstance() || obj->IsInstance(); } - }); + return ffi::WalkResult::Advance(); + }; + ffi::StructuralWalk(e, walk_fn); if (is_simple && !cand_vars.empty()) { for (const PrimVar& new_var : cand_vars) { if (!std::any_of(vars.begin(), vars.end(), diff --git a/tests/cpp/ir_functor_test.cc b/tests/cpp/ir_functor_test.cc index c197b4931c38..192f7f5abda6 100644 --- a/tests/cpp/ir_functor_test.cc +++ b/tests/cpp/ir_functor_test.cc @@ -34,6 +34,7 @@ #include #include +#include TEST(IRF, Basic) { using namespace tvm; @@ -55,13 +56,18 @@ TEST(IRF, CountVar) { PrimVar x("x"), y("y"); auto z = x + 1 + y + y; - tirx::PostOrderVisit(z, [&n_var](const ffi::ObjectRef& n) { - if (n.as()) ++n_var; - }); + std::unordered_set visited; + auto walk_fn = [&n_var, &visited](const Var& var) -> ffi::Expected { + if (visited.insert(var.get()).second) { + ++n_var; + } + return ffi::WalkResult::Advance(); + }; + ffi::StructuralWalk(z, walk_fn); TVM_FFI_ICHECK_EQ(n_var, 2); } -TEST(IRF, PreOrderVisit) { +TEST(IRF, PreOrderStructuralWalk) { using namespace tvm; using namespace tvm::tirx; Stmt init = @@ -73,24 +79,23 @@ TEST(IRF, PreOrderVisit) { bool init_visited = false; bool stopped_at_if = true; bool body_visited = false; - PreOrderVisit(block, [&](const ffi::ObjectRef& n) -> bool { - if (n->IsInstance()) { - init_visited = true; - return false; - } - if (const auto* eval = n.as()) { - if (const auto* int_imm = eval->value.as()) { - if (int_imm->value == 0) { - stopped_at_if = false; - } else if (int_imm->value == 1) { - body_visited = true; - } else { - TVM_FFI_THROW(InternalError) << "Unreachable"; - } + auto visit_if = [&](const IfThenElse&) -> ffi::Expected { + init_visited = true; + return ffi::WalkResult::Skip(); + }; + auto visit_evaluate = [&](const Evaluate& eval) -> ffi::Expected { + if (const auto* int_imm = eval->value.as()) { + if (int_imm->value == 0) { + stopped_at_if = false; + } else if (int_imm->value == 1) { + body_visited = true; + } else { + TVM_FFI_THROW(InternalError) << "Unreachable"; } } - return true; - }); + return ffi::WalkResult::Advance(); + }; + ffi::StructuralWalk(block, visit_if, visit_evaluate); ASSERT_EQ(init_visited, true); ASSERT_EQ(stopped_at_if, true); ASSERT_EQ(body_visited, true); @@ -366,7 +371,8 @@ TEST(IRF, StructuralMapSplicesMappedSeqStmtChild) { { Stmt input = make_input(); Stmt shared = input; - Stmt mapped = ffi::StructuralMap(input, expand_one).cast(); + Stmt mapped = + ffi::StructuralMap(input, expand_one).as_or_throw(); EXPECT_FALSE(mapped.same_as(input)); EXPECT_EQ(shared.as()->seq.size(), 3); check_values(mapped, {5, 2, 3, 4}); @@ -376,8 +382,8 @@ TEST(IRF, StructuralMapSplicesMappedSeqStmtChild) { Stmt input = make_input(); const auto* original = input.get(); const auto* original_array = input.as()->seq.GetArrayObj(); - Stmt mapped = - ffi::StructuralMap(std::move(input), expand_one).cast(); + Stmt mapped = ffi::StructuralMap(std::move(input), expand_one) + .as_or_throw(); EXPECT_EQ(mapped.get(), original); EXPECT_NE(mapped.as()->seq.GetArrayObj(), original_array); check_values(mapped, {5, 2, 3, 4}); @@ -391,14 +397,14 @@ TEST(IRF, StructuralMapSplicesMappedSeqStmtChild) { bool expect_array_reuse) { Stmt ordinary_input = make_boundary_input(); Stmt ordinary = ffi::StructuralMap(ordinary_input, transform) - .template cast(); + .template as_or_throw(); Stmt inplace_input = make_boundary_input(); const auto* original_root = inplace_input.get(); const auto* original_array = inplace_input.as()->seq.GetArrayObj(); Stmt inplace = ffi::StructuralMap(std::move(inplace_input), transform) - .template cast(); + .template as_or_throw(); EXPECT_EQ(inplace.get(), original_root); if (expect_array_reuse) { @@ -477,12 +483,12 @@ TEST(IRF, StructuralMapSplicesMappedSeqStmtChild) { auto remove_all = [](const Evaluate&) -> Stmt { return Evaluate(0); }; Stmt ordinary_input = make_boundary_input(); - Stmt ordinary = - ffi::StructuralMap(ordinary_input, remove_all).cast(); + Stmt ordinary = ffi::StructuralMap(ordinary_input, remove_all) + .as_or_throw(); Stmt inplace_input = make_boundary_input(); Stmt inplace = ffi::StructuralMap(std::move(inplace_input), remove_all) - .cast(); + .as_or_throw(); EXPECT_TRUE(ffi::StructuralEqual()(ordinary, inplace)); for (const Stmt& result : {ordinary, inplace}) { const auto* evaluate = result.as(); @@ -497,10 +503,11 @@ TEST(IRF, StructuralMapSplicesMappedSeqStmtChild) { return value != nullptr && value->value == 4 ? Stmt(evaluate) : Stmt(Evaluate(0)); }; ordinary_input = make_boundary_input(); - ordinary = ffi::StructuralMap(ordinary_input, keep_last).cast(); + ordinary = + ffi::StructuralMap(ordinary_input, keep_last).as_or_throw(); inplace_input = make_boundary_input(); inplace = ffi::StructuralMap(std::move(inplace_input), keep_last) - .cast(); + .as_or_throw(); EXPECT_TRUE(ffi::StructuralEqual()(ordinary, inplace)); for (const Stmt& result : {ordinary, inplace}) { const auto* evaluate = result.as(); @@ -523,8 +530,8 @@ TEST(IRF, StructuralMapPreservesSeqStmtElementUniqueness) { Stmt input = SeqStmt({Evaluate(IntImm::Int32(1)), Evaluate(IntImm::Int32(3))}); const auto* original_root = input.get(); const auto* original_first = input.as()->seq[0].as(); - Stmt mapped = - ffi::StructuralMap(std::move(input), replace_one).cast(); + Stmt mapped = ffi::StructuralMap(std::move(input), replace_one) + .as_or_throw(); const auto* mapped_seq = mapped.as(); ASSERT_NE(mapped_seq, nullptr); @@ -537,8 +544,8 @@ TEST(IRF, StructuralMapPreservesSeqStmtElementUniqueness) { ffi::Array shared_seq = {Evaluate(IntImm::Int32(1)), Evaluate(IntImm::Int32(3))}; const auto* shared_first = shared_seq[0].as(); Stmt input = SeqStmt(shared_seq); - Stmt mapped = - ffi::StructuralMap(std::move(input), replace_one).cast(); + Stmt mapped = ffi::StructuralMap(std::move(input), replace_one) + .as_or_throw(); const auto* mapped_seq = mapped.as(); ASSERT_NE(mapped_seq, nullptr); @@ -552,7 +559,7 @@ TEST(IRF, StructuralMapPreservesSeqStmtElementUniqueness) { auto no_float_match = [](const FloatImm& value) -> PrimExpr { return value; }; Stmt unchanged = ffi::StructuralMap(std::move(unchanged_input), no_float_match) - .cast(); + .as_or_throw(); EXPECT_EQ(unchanged.get(), unchanged_root); EXPECT_TRUE(unchanged.as()->seq.same_as(shared_seq)); } @@ -577,14 +584,13 @@ TEST(IRF, StructuralHooksPreserveScopeIdDefRegions) { int binder = -1; int extent = -1; int preferred = -1; - ffi::StructuralWalk( - make_input(), - [&](const Var& var, TVMFFIDefRegionKind kind) -> ffi::Expected { - if (var->name == "binder") binder = kind; - if (var->name == "extent") extent = kind; - if (var->name == "preferred") preferred = kind; - return ffi::WalkResult::Advance(); - }); + auto walk_fn = [&](const Var& var, TVMFFIDefRegionKind kind) -> ffi::Expected { + if (var->name == "binder") binder = kind; + if (var->name == "extent") extent = kind; + if (var->name == "preferred") preferred = kind; + return ffi::WalkResult::Advance(); + }; + ffi::StructuralWalk(make_input(), walk_fn); check_kinds(binder, extent, preferred); } diff --git a/tests/python/relax/test_backend_dispatch_sort_scan.py b/tests/python/relax/test_backend_dispatch_sort_scan.py index 7fcd26491759..df4aca033a42 100644 --- a/tests/python/relax/test_backend_dispatch_sort_scan.py +++ b/tests/python/relax/test_backend_dispatch_sort_scan.py @@ -18,6 +18,7 @@ import numpy as np import pytest +import tvm_ffi import tvm import tvm.script @@ -542,7 +543,7 @@ def collect_floor_divisors(node): if isinstance(node, tirx.FloorDiv): floor_divisors.append(node.b) - tirx.stmt_functor.post_order_visit(cumsum.body, collect_floor_divisors) + tvm_ffi.structural_walk(cumsum.body, collect_floor_divisors) assert floor_divisors assert all( isinstance(divisor, tirx.IntImm) diff --git a/tests/python/relax/test_pipeline.py b/tests/python/relax/test_pipeline.py index ae7a889940f1..dd8a15f8194f 100644 --- a/tests/python/relax/test_pipeline.py +++ b/tests/python/relax/test_pipeline.py @@ -176,7 +176,7 @@ def _visit(node): if isinstance(node, tvm.tirx.For) and node.kind == tvm.tirx.ForKind.THREAD_BINDING: found = True - tvm.tirx.stmt_functor.post_order_visit(func.body, _visit) + tvm_ffi.structural_walk(func.body, _visit) return found diff --git a/tests/python/s_tir/base/test_sblock_dependence_info.py b/tests/python/s_tir/base/test_sblock_dependence_info.py index eb6ee6841d5f..6d20e16d4928 100644 --- a/tests/python/s_tir/base/test_sblock_dependence_info.py +++ b/tests/python/s_tir/base/test_sblock_dependence_info.py @@ -20,6 +20,7 @@ import sys import pytest +import tvm_ffi import tvm import tvm.testing @@ -29,7 +30,6 @@ from tvm.s_tir.sblock_scope import DepKind from tvm.script import tirx as T from tvm.tirx import PrimFunc -from tvm.tirx.stmt_functor import post_order_visit # pylint: disable=no-member,invalid-name,unused-variable @@ -94,7 +94,7 @@ def update_blocks(node): blocks[node.name_hint] = node # post_order_visit(func.body, lambda node: blocks[node.name_hint] = node if isinstance(node, tvm.tirx.SBlock) else None) - post_order_visit(func.body, update_blocks) + tvm_ffi.structural_walk(func.body, update_blocks) return blocks diff --git a/tests/python/s_tir/schedule/test_tir_schedule_analysis.py b/tests/python/s_tir/schedule/test_tir_schedule_analysis.py index 48e7097b1b9d..3fb208e1afa1 100644 --- a/tests/python/s_tir/schedule/test_tir_schedule_analysis.py +++ b/tests/python/s_tir/schedule/test_tir_schedule_analysis.py @@ -18,6 +18,7 @@ # ruff: noqa: F401, F841 import pytest +import tvm_ffi import tvm import tvm.testing @@ -49,7 +50,6 @@ ) from tvm.tirx.analysis import expr_deep_equal from tvm.tirx.function import TensorIntrin -from tvm.tirx.stmt_functor import pre_order_visit def _make_vars(*args: str) -> list[Var]: @@ -227,9 +227,9 @@ def collect_loops(prim_func): def callback(node): if isinstance(node, tvm.tirx.For): loops.append(node) - return True + return tvm_ffi.WalkResult.ADVANCE - pre_order_visit(prim_func.body, callback) + tvm_ffi.structural_walk(prim_func.body, (object, callback), order="pre") return loops diff --git a/tests/python/s_tir/schedule/test_tir_schedule_block_scope.py b/tests/python/s_tir/schedule/test_tir_schedule_block_scope.py index f98b45c4ec98..f18e49cc48d4 100644 --- a/tests/python/s_tir/schedule/test_tir_schedule_block_scope.py +++ b/tests/python/s_tir/schedule/test_tir_schedule_block_scope.py @@ -19,13 +19,13 @@ import sys import pytest +from tvm_ffi import structural_walk as post_order_visit import tvm import tvm.testing from tvm import s_tir, tirx from tvm.s_tir.schedule import DepKind from tvm.script import tirx as T -from tvm.tirx.stmt_functor import post_order_visit # pylint: disable=no-member,invalid-name,unused-variable diff --git a/tests/python/s_tir/schedule/test_tir_schedule_state_cached_flags.py b/tests/python/s_tir/schedule/test_tir_schedule_state_cached_flags.py index 02f56a156b2c..cccc4fd38dc0 100644 --- a/tests/python/s_tir/schedule/test_tir_schedule_state_cached_flags.py +++ b/tests/python/s_tir/schedule/test_tir_schedule_state_cached_flags.py @@ -19,13 +19,13 @@ import sys import pytest +from tvm_ffi import structural_walk as post_order_visit import tvm import tvm.testing from tvm import s_tir, tirx from tvm.s_tir.schedule.state import CachedFlags from tvm.script import tirx as T -from tvm.tirx.stmt_functor import post_order_visit # pylint: disable=no-member,invalid-name,unused-variable,unexpected-keyword-arg # fmt: off diff --git a/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py b/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py index 4252f88f72b4..9cd1907cff1e 100644 --- a/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py +++ b/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py @@ -19,6 +19,7 @@ import sys import pytest +import tvm_ffi import tvm import tvm.testing @@ -404,6 +405,30 @@ def test_transform_block_layout_fail_mixed_iter_type(use_block_name): ) +def test_mixed_iter_type_detection_interrupts_walk(): + spatial = tirx.Var("spatial", "int32") + reduction = tirx.Var("reduction", "int32") + unreachable = tirx.Var("unreachable", "int32") + visited = [] + + def detect(var): + visited.append(var) + if var.same_as(reduction): + return tvm_ffi.VisitInterrupt() + return None + + result = tvm_ffi.structural_walk( + (spatial + reduction) + unreachable, + (tirx.Var, detect), + order="post", + ) + + assert isinstance(result, tvm_ffi.VisitInterrupt) + assert any(var.same_as(spatial) for var in visited) + assert any(var.same_as(reduction) for var in visited) + assert not any(var.same_as(unreachable) for var in visited) + + def test_transform_block_layout_int64_extent(use_block_name): @T.prim_func(s_tir=True) def elementwise_int64_extent( diff --git a/tests/python/s_tir/transform/test_s_tir_transform_hoist_if.py b/tests/python/s_tir/transform/test_s_tir_transform_hoist_if.py index d59aedb4d24c..8915aeebc83e 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_hoist_if.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_hoist_if.py @@ -41,7 +41,7 @@ def _visit(op): key = op if isinstance(op, tvm.tirx.IfThenElse): global var_list - tvm.tirx.stmt_functor.post_order_visit(op.condition, _extract_vars) + tvm_ffi.structural_walk(op.condition, _extract_vars) val = [(op.then_case, op.else_case), ("tirx.IfThenElse", tuple(var_list))] var_list.clear() elif isinstance(op, tvm.tirx.For): @@ -52,7 +52,7 @@ def _visit(op): return node_dict[key] = val - tvm.tirx.stmt_functor.post_order_visit(stmt, _visit) + tvm_ffi.structural_walk(stmt, _visit) for key, val in node_dict.items(): struct[val[1]] = tuple( node_dict[child][1] if child in node_dict else None for child in val[0] diff --git a/tests/python/s_tir/transform/test_s_tir_transform_inject_double_buffer.py b/tests/python/s_tir/transform/test_s_tir_transform_inject_double_buffer.py index 5a761178b3c3..d8824ca6fcda 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_inject_double_buffer.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_inject_double_buffer.py @@ -16,6 +16,8 @@ # under the License. # ruff: noqa: F841 +import tvm_ffi + import tvm import tvm.testing from tvm.script import ir as I @@ -59,7 +61,7 @@ def visitor(op): if isinstance(op, tvm.tirx.AllocBuffer) and "B" in str(op.buffer.data): allocate_node = op - tvm.tirx.stmt_functor.post_order_visit(stmt, visitor) + tvm_ffi.structural_walk(stmt, visitor) assert allocate_node is not None assert list(allocate_node.buffer.shape) == [m * 2] @@ -70,7 +72,7 @@ def count_sync(op): if isinstance(op, tvm.ir.Call) and op.op.same_as(tvm.ir.Op.get("tirx.tvm_storage_sync")): count[0] += 1 - tvm.tirx.stmt_functor.post_order_visit(f.body, count_sync) + tvm_ffi.structural_walk(f.body, count_sync) assert count[0] == 4 @@ -107,7 +109,7 @@ def visitor(op): if isinstance(op, tvm.tirx.AllocBuffer): allocate_node = op - tvm.tirx.stmt_functor.post_order_visit(After["main"].body, visitor) + tvm_ffi.structural_walk(After["main"].body, visitor) assert allocate_node is not None assert list(allocate_node.buffer.shape) == [64] diff --git a/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py b/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py index a53f3c439c12..922bcb992f5d 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py @@ -50,7 +50,7 @@ def verify(n): if isinstance(n, tvm.ir.Call) and n.op.name == "tirx.s_tir.cp_async_raw": num_alloc[0] += 1 - tvm.tirx.stmt_functor.post_order_visit(stmt, verify) + tvm_ffi.structural_walk(stmt, verify) return num_alloc[0] diff --git a/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_ldg32.py b/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_ldg32.py index 67e568a248f1..2c190279414f 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_ldg32.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_ldg32.py @@ -16,6 +16,8 @@ # under the License. # ruff: noqa: F401 +import tvm_ffi + import tvm import tvm.testing from tvm import s_tir @@ -29,7 +31,7 @@ def visit(n): if isinstance(n, tvm.tirx.AllocBuffer): num_alloc[0] += 1 - tvm.tirx.stmt_functor.post_order_visit(stmt, visit) + tvm_ffi.structural_walk(stmt, visit) return num_alloc[0] @@ -40,7 +42,7 @@ def visit(n): if isinstance(n, tvm.ir.Call) and n.op.name == "tirx.s_tir.ldg32": num_call[0] += 1 - tvm.tirx.stmt_functor.post_order_visit(stmt, visit) + tvm_ffi.structural_walk(stmt, visit) return num_call[0] diff --git a/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py b/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py index 471974f16c58..315ec04de657 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py @@ -15,6 +15,8 @@ # specific language governing permissions and limitations # under the License. # ruff: noqa: F841 +import tvm_ffi + import tvm import tvm.testing from tvm.script import ir as I @@ -60,7 +62,7 @@ def find_allocates(node): if isinstance(node, tvm.tirx.AllocBuffer): allocates.append(node) - tvm.tirx.stmt_functor.post_order_visit(stmt.body, find_allocates) + tvm_ffi.structural_walk(stmt.body, find_allocates) assert len(allocates) == 1 assert list(allocates[0].buffer.ty.shape) == [B_expected_alloc] @@ -109,7 +111,7 @@ def find_allocates(node): if isinstance(node, tvm.tirx.AllocBuffer): allocates.append(node) - tvm.tirx.stmt_functor.post_order_visit(stmt.body, find_allocates) + tvm_ffi.structural_walk(stmt.body, find_allocates) assert len(allocates) == 3 # Check that we have the expected extents (order may vary) extents = sorted([int(a.buffer.ty.shape[0]) for a in allocates]) @@ -145,7 +147,7 @@ def find_ifs(node): if isinstance(node, tvm.tirx.IfThenElse): if_nodes.append(node) - tvm.tirx.stmt_functor.post_order_visit(stmt.body, find_ifs) + tvm_ffi.structural_walk(stmt.body, find_ifs) assert len(if_nodes) == 2 # First if has else_case, second does not @@ -207,7 +209,7 @@ def visitor(op): if isinstance(op, tvm.tirx.AllocBuffer) and "shared" in str(op.buffer.data.ty): allocate_node = op - tvm.tirx.stmt_functor.post_order_visit(after_func.body, visitor) + tvm_ffi.structural_walk(after_func.body, visitor) assert allocate_node is not None assert list(allocate_node.buffer.ty.shape) == [4] assert allocate_node.buffer.ty.dtype == "int32x4" @@ -236,7 +238,7 @@ def visitor(node): }: masked_calls.append(node) - tvm.tirx.stmt_functor.post_order_visit(after.body, visitor) + tvm_ffi.structural_walk(after.body, visitor) assert len(masked_calls) == 4 assert all(list(call.args[0].ty.shape) == [8] for call in masked_calls) analyzer = tvm.arith.Analyzer() diff --git a/tests/python/s_tir/transform/test_s_tir_transform_loop_partition.py b/tests/python/s_tir/transform/test_s_tir_transform_loop_partition.py index aa111bed1dca..be1cfe5f906f 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_loop_partition.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_loop_partition.py @@ -17,6 +17,7 @@ # ruff: noqa: F401 import numpy import pytest +import tvm_ffi import tvm import tvm.testing @@ -26,7 +27,7 @@ def collect_visit(stmt, f): ret = [] - tvm.tirx.stmt_functor.post_order_visit(stmt, lambda x: ret.append(f(x))) + tvm_ffi.structural_walk(stmt, lambda x: ret.append(f(x))) return ret diff --git a/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_all_reduce.py b/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_all_reduce.py index 5c08e66a9556..17253785e1a4 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_all_reduce.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_all_reduce.py @@ -16,6 +16,8 @@ # under the License. # ruff: noqa: F401, F841 +import tvm_ffi + import tvm import tvm.testing from tvm import s_tir @@ -31,7 +33,7 @@ def visit(node): if isinstance(node, tvm.tirx.AllocBuffer) and "tirx.volatile" in node.annotations: has_volatile_alloc = has_volatile_alloc or node.annotations["tirx.volatile"] is True - tvm.tirx.stmt_functor.post_order_visit(mod["main"].body, visit) + tvm_ffi.structural_walk(mod["main"].body, visit) return has_volatile_alloc diff --git a/tests/python/s_tir/transform/test_s_tir_transform_memhammer_lower_auto_copy.py b/tests/python/s_tir/transform/test_s_tir_transform_memhammer_lower_auto_copy.py index 89a66cd4cc65..5d050149f8ef 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_memhammer_lower_auto_copy.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_memhammer_lower_auto_copy.py @@ -19,6 +19,7 @@ import sys import pytest +import tvm_ffi import tvm from tvm import s_tir @@ -1146,7 +1147,7 @@ def verify(n): for buf in n.alloc_buffers: alloc_extents.append(buf.shape) - tvm.tirx.stmt_functor.post_order_visit(stmt, verify) + tvm_ffi.structural_walk(stmt, verify) assert num_alloc[0] == 1 if alloc_size: diff --git a/tests/python/tirx-base/test_tir_base.py b/tests/python/tirx-base/test_tir_base.py index dbf1eb623fc8..c82aad76bc72 100644 --- a/tests/python/tirx-base/test_tir_base.py +++ b/tests/python/tirx-base/test_tir_base.py @@ -134,7 +134,7 @@ def test_return_stmt_functor_traversal_and_mutation(): stmt = tirx.Return(x + 1, span) visited = [] - tirx.stmt_functor.post_order_visit(stmt, visited.append) + tvm_ffi.structural_walk(stmt, visited.append) assert any(node.same_as(x) for node in visited) assert any(isinstance(node, tirx.Return) for node in visited) diff --git a/tests/python/tirx-transform/test_tir_transform_bf16_legalize.py b/tests/python/tirx-transform/test_tir_transform_bf16_legalize.py index ce1437859843..0a23e75b2e14 100644 --- a/tests/python/tirx-transform/test_tir_transform_bf16_legalize.py +++ b/tests/python/tirx-transform/test_tir_transform_bf16_legalize.py @@ -14,6 +14,8 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +import tvm_ffi + import tvm import tvm.script from tvm.script import tirx as T @@ -147,7 +149,16 @@ def main(Aptr: T.handle("bfloat16"), Cptr: T.handle("bfloat16")): def collect(mod): nodes = [] - tvm.tirx.stmt_functor.post_order_visit(mod["main"].body, nodes.append) + seen = [] + + def append_once(node): + if isinstance(node, tvm_ffi.Object) and not any( + node.same_as(previous) for previous in seen + ): + seen.append(node) + nodes.append(node) + + tvm_ffi.structural_walk(mod["main"].body, append_once) buffers = { node.buffer.name: str(node.buffer.dtype) for node in nodes diff --git a/tests/python/tirx-transform/test_tir_transform_lower_intrin.py b/tests/python/tirx-transform/test_tir_transform_lower_intrin.py index 052e2522f19e..86d986250e26 100644 --- a/tests/python/tirx-transform/test_tir_transform_lower_intrin.py +++ b/tests/python/tirx-transform/test_tir_transform_lower_intrin.py @@ -17,6 +17,7 @@ # ruff: noqa: RUF005 import numpy as np import pytest +import tvm_ffi import tvm import tvm.testing @@ -106,7 +107,7 @@ def collect(node): elif node.op.name == "tirx.address_of": address_calls.append(node) - tvm.tirx.stmt_functor.post_order_visit(lowered.body, collect) + tvm_ffi.structural_walk(lowered.body, collect) assert not access_ptr_calls assert len(address_calls) == 1 load = address_calls[0].args[0] diff --git a/tests/python/tirx-transform/test_tir_transform_make_packed_api.py b/tests/python/tirx-transform/test_tir_transform_make_packed_api.py index 81d46061ff9d..e48c5aac780b 100644 --- a/tests/python/tirx-transform/test_tir_transform_make_packed_api.py +++ b/tests/python/tirx-transform/test_tir_transform_make_packed_api.py @@ -21,6 +21,7 @@ """ import pytest +import tvm_ffi import tvm import tvm.testing @@ -37,7 +38,7 @@ def _visitor(stmt): nonlocal result result = stmt - tirx.stmt_functor.post_order_visit(func.body, _visitor) + tvm_ffi.structural_walk(func.body, _visitor) return result @@ -215,7 +216,7 @@ def collect(node): if isinstance(field, tvm.tirx.IntImm) and int(field) == 13: return_type_indices.append(int(node.args[3])) - tvm.tirx.stmt_functor.post_order_visit(after.body, collect) + tvm_ffi.structural_walk(after.body, collect) assert 4 in return_type_indices # ffi::TypeIndex::kTVMFFIOpaquePtr @@ -482,7 +483,7 @@ def collect(node): if isinstance(node, tirx.DeclBuffer): declared_buffers.append(node.buffer) - tirx.stmt_functor.post_order_visit(after.body, collect) + tvm_ffi.structural_walk(after.body, collect) assert len(alignment_nodes) == 1 assert any(alignment_nodes[0].same_as(buffer) for buffer in declared_buffers) diff --git a/tests/python/tirx-transform/test_tir_transform_pointer_value_type_rewrite.py b/tests/python/tirx-transform/test_tir_transform_pointer_value_type_rewrite.py index 89fd4011a6d6..e8a387cee3b2 100644 --- a/tests/python/tirx-transform/test_tir_transform_pointer_value_type_rewrite.py +++ b/tests/python/tirx-transform/test_tir_transform_pointer_value_type_rewrite.py @@ -16,6 +16,8 @@ # under the License. # pylint: disable=invalid-name, missing-docstring +import tvm_ffi + import tvm import tvm.testing from tvm.script import ir as I @@ -161,7 +163,7 @@ def main(A: T.Buffer((16,), "float32")): decl_buffers = [] buffer_stores = [] - tvm.tirx.stmt_functor.post_order_visit( + tvm_ffi.structural_walk( func.body, lambda node: ( decl_buffers.append(node) diff --git a/tests/python/tirx-transform/test_tir_transform_split_host_device.py b/tests/python/tirx-transform/test_tir_transform_split_host_device.py index 2990efc72504..ed5ce0f9a32b 100644 --- a/tests/python/tirx-transform/test_tir_transform_split_host_device.py +++ b/tests/python/tirx-transform/test_tir_transform_split_host_device.py @@ -16,6 +16,7 @@ # under the License. import pytest +import tvm_ffi import tvm import tvm.testing @@ -362,7 +363,7 @@ def collect(node): if isinstance(node, tvm.tirx.DeclBuffer): declared_buffers.append(node.buffer) - tvm.tirx.stmt_functor.post_order_visit(kernel.body, collect) + tvm_ffi.structural_walk(kernel.body, collect) assert len(declared_buffers) == 1 assert not tvm.tirx.analysis.undefined_vars(kernel.body, kernel.params) diff --git a/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py b/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py index d5b271b4689f..b3133cba1e91 100644 --- a/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py +++ b/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py @@ -18,6 +18,7 @@ import sys import pytest +import tvm_ffi import tvm import tvm.testing @@ -48,7 +49,7 @@ def verify(n): num_alloc[0] += 1 assert n.buffer.ty.shape[0].value == 200 - tvm.tirx.stmt_functor.post_order_visit(body, verify) + tvm_ffi.structural_walk(body, verify) assert num_alloc[0] == 1 @@ -106,7 +107,7 @@ def verify(n): offset = offset_generater(dtype_list, length) body = tvm.tirx.transform.StorageRewrite()(mod)["func"].body - tvm.tirx.stmt_functor.post_order_visit(body, verify) + tvm_ffi.structural_walk(body, verify) length = 1024 dtype_list = ["float16", "int32", "uint16", "int8"] @@ -161,12 +162,12 @@ def verify(n): total_alloc = [0] mod = tvm.IRModule.from_expr(before.with_attr("global_symbol", "main")) - tvm.tirx.stmt_functor.post_order_visit(mod["main"].body, verify) + tvm_ffi.structural_walk(mod["main"].body, verify) assert total_alloc[0] == 24 total_alloc[0] = 0 mod = tvm.tirx.transform.StorageRewrite()(mod) - tvm.tirx.stmt_functor.post_order_visit(mod["main"].body, verify) + tvm_ffi.structural_walk(mod["main"].body, verify) assert total_alloc[0] == 16 @@ -244,13 +245,13 @@ def count_alloc(n): if isinstance(n, tvm.tirx.AllocBuffer): num_alloc[0] += 1 - tvm.tirx.stmt_functor.post_order_visit(inner, count_alloc) + tvm_ffi.structural_walk(inner, count_alloc) assert num_alloc[0] == 2 # j and A allocations mod = tvm.IRModule.from_expr(func_serial) body = tvm.tirx.transform.StorageRewrite()(mod)["func_serial"] num_alloc[0] = 0 - tvm.tirx.stmt_functor.post_order_visit(body.body, count_alloc) + tvm_ffi.structural_walk(body.body, count_alloc) assert num_alloc[0] == 2 # j and A allocations @@ -282,7 +283,7 @@ def verify(n): num_alloc[0] += 1 assert n.buffer.ty.shape[0].value == 500 - tvm.tirx.stmt_functor.post_order_visit(body, verify) + tvm_ffi.structural_walk(body, verify) assert num_alloc[0] == 1 @@ -312,7 +313,7 @@ def verify(n): num_alloc[0] += 1 assert n.buffer.ty.shape[0].value == 200 - tvm.tirx.stmt_functor.post_order_visit(body, verify) + tvm_ffi.structural_walk(body, verify) assert num_alloc[0] == 1 @@ -344,7 +345,7 @@ def verify(n): num_alloc[0] += 1 assert n.buffer.ty.shape[0].value == 800 - tvm.tirx.stmt_functor.post_order_visit(body, verify) + tvm_ffi.structural_walk(body, verify) assert num_alloc[0] == 1 @@ -475,7 +476,7 @@ def func(D: T.Buffer(1, "float32")): after = tvm.tirx.transform.StorageRewrite()(tvm.IRModule.from_expr(func))["func"] allocations = [] - tvm.tirx.stmt_functor.post_order_visit( + tvm_ffi.structural_walk( after.body, lambda node: allocations.append(node) if isinstance(node, tvm.tirx.AllocBuffer) else None, ) diff --git a/tests/python/tirx-transform/test_tir_transform_vectorize.py b/tests/python/tirx-transform/test_tir_transform_vectorize.py index e7ac072cb787..e0940db94a19 100644 --- a/tests/python/tirx-transform/test_tir_transform_vectorize.py +++ b/tests/python/tirx-transform/test_tir_transform_vectorize.py @@ -16,6 +16,7 @@ # under the License. # ruff: noqa: F841 import pytest +import tvm_ffi import tvm import tvm.testing @@ -664,7 +665,7 @@ def collect_predicates(node): if isinstance(node, tvm.ir.Call) and node.op.name == "tirx.masked_store": predicates.append(node.args[-1]) - tvm.tirx.stmt_functor.post_order_visit(after.body, collect_predicates) + tvm_ffi.structural_walk(after.body, collect_predicates) assert len(predicates) == 2 assert any( isinstance(predicate, tvm.ir.Call) and predicate.op.name == "ir.prim.bitwise_and" diff --git a/tests/python/tirx/codegen/test_codegen_blackwell.py b/tests/python/tirx/codegen/test_codegen_blackwell.py index f02103a20b33..26138ffd685d 100644 --- a/tests/python/tirx/codegen/test_codegen_blackwell.py +++ b/tests/python/tirx/codegen/test_codegen_blackwell.py @@ -17,6 +17,7 @@ # pylint: disable=missing-function-docstring import numpy as np import pytest +import tvm_ffi import tvm import tvm.testing @@ -53,7 +54,7 @@ def visit(node): if isinstance(node, tvm.ir.Call) and node.op.name == arrive_op_name: arrive_calls.append(node) - tvm.tirx.stmt_functor.post_order_visit(func.body, visit) + tvm_ffi.structural_walk(func.body, visit) assert len(bindings) == 1 assert len(buffers) == 1 assert len(mapa_calls) == 1 @@ -201,7 +202,7 @@ def visit(node): if isinstance(node, tvm.ir.Call) and node.op.name == "tirx.ptx.mbarrier_arrive": arrive_calls.append(node) - tvm.tirx.stmt_functor.post_order_visit(test_local_arrive.body, visit) + tvm_ffi.structural_walk(test_local_arrive.body, visit) assert len(arrive_calls) == 1 call = arrive_calls[0] assert call.args[1].value == 2 diff --git a/tests/python/tirx/codegen/test_codegen_dsmem.py b/tests/python/tirx/codegen/test_codegen_dsmem.py index d3ecd16a6202..19481a5e3ca2 100644 --- a/tests/python/tirx/codegen/test_codegen_dsmem.py +++ b/tests/python/tirx/codegen/test_codegen_dsmem.py @@ -17,6 +17,8 @@ # pylint: disable=missing-function-docstring """Tests for cp.async.bulk.shared::cluster.shared::cta PTX instruction codegen.""" +import tvm_ffi + import tvm import tvm.testing from tvm.ir import PointerType, PrimType, assert_structural_equal @@ -123,7 +125,7 @@ def collect(node): elif isinstance(node, tvm.ir.TensorLoad): loads.append(node) - tvm.tirx.stmt_functor.post_order_visit(main.body, collect) + tvm_ffi.structural_walk(main.body, collect) assert len(binds) == 1 assert isinstance(binds[0].var.ty, PointerType) assert binds[0].var.ty.storage_scope == "shared" diff --git a/tests/python/tirx/codegen/test_ptx_addr.py b/tests/python/tirx/codegen/test_ptx_addr.py index e333f9bb614d..b26750370d1e 100644 --- a/tests/python/tirx/codegen/test_ptx_addr.py +++ b/tests/python/tirx/codegen/test_ptx_addr.py @@ -17,6 +17,7 @@ """Tests for ``T.ptx.addr(base, byte_offset)``.""" import pytest +import tvm_ffi import tvm from tvm.ir import Call, Op @@ -40,7 +41,7 @@ def visit(node): if isinstance(node, Call) and getattr(node.op, "name", None) == op_name: calls.append(node) - tvm.tirx.stmt_functor.post_order_visit(func.body, visit) + tvm_ffi.structural_walk(func.body, visit) return calls diff --git a/tests/python/tirx/test_hint.py b/tests/python/tirx/test_hint.py index 1b3ca5648784..6186c471e35b 100644 --- a/tests/python/tirx/test_hint.py +++ b/tests/python/tirx/test_hint.py @@ -16,6 +16,8 @@ # under the License. """Tests for T.hint() — universal directive primitive for TIRx sketch language.""" +import tvm_ffi + import tvm import tvm.script import tvm.testing @@ -50,7 +52,7 @@ def visit(stmt): assert str(stmt.node["message"]) == "persistent tile scheduler with L2 swizzle" found[0] = True - tvm.tirx.stmt_functor.post_order_visit(func.body, visit) + tvm_ffi.structural_walk(func.body, visit) assert found[0], "Expected AttrStmt with attr_key='tirx_hint' not found" @@ -74,7 +76,7 @@ def visit(stmt): assert str(stmt.node["message"]) == "software pipeline, depth 4" found[0] = True - tvm.tirx.stmt_functor.post_order_visit(func.body, visit) + tvm_ffi.structural_walk(func.body, visit) assert found[0], "Expected AttrStmt with attr_key='tirx_hint' not found" @@ -100,7 +102,7 @@ def visit(stmt): assert str(stmt.node["depth"]) == "4" found[0] = True - tvm.tirx.stmt_functor.post_order_visit(func.body, visit) + tvm_ffi.structural_walk(func.body, visit) assert found[0], "Expected AttrStmt with attr_key='tirx_hint' not found" @@ -220,7 +222,7 @@ def visit(stmt): assert isinstance(stmt.node["access"], BufferRegion) found[0] = True - tvm.tirx.stmt_functor.post_order_visit(func.body, visit) + tvm_ffi.structural_walk(func.body, visit) assert found[0], "Expected AttrStmt with attr_key='tirx_hint' containing access not found" @@ -251,7 +253,7 @@ def visit(stmt): assert len(br.region) == 2 found[0] = True - tvm.tirx.stmt_functor.post_order_visit(func.body, visit) + tvm_ffi.structural_walk(func.body, visit) assert found[0], "Expected AttrStmt with structured BufferRegion access not found" diff --git a/tests/python/tirx/test_op_namespace_cleanup.py b/tests/python/tirx/test_op_namespace_cleanup.py index 09ad39d4a676..78f56aa3d920 100644 --- a/tests/python/tirx/test_op_namespace_cleanup.py +++ b/tests/python/tirx/test_op_namespace_cleanup.py @@ -21,6 +21,7 @@ import types import pytest +import tvm_ffi import tvm from tvm.ir import Op, assert_structural_equal @@ -36,7 +37,7 @@ def visit(stmt): if isinstance(stmt, TilePrimitiveCall): calls.append(stmt) - tvm.tirx.stmt_functor.post_order_visit(func.body, visit) + tvm_ffi.structural_walk(func.body, visit) return calls @@ -47,7 +48,7 @@ def visit(node): if isinstance(node, tvm.ir.Call): calls.append(node) - tvm.tirx.stmt_functor.post_order_visit(func.body, visit) + tvm_ffi.structural_walk(func.body, visit) return calls diff --git a/tests/python/tirx/test_parser_printer.py b/tests/python/tirx/test_parser_printer.py index c6f6d5c21ede..45ebe11b13f5 100644 --- a/tests/python/tirx/test_parser_printer.py +++ b/tests/python/tirx/test_parser_printer.py @@ -17,6 +17,7 @@ import math import pytest +import tvm_ffi import tvm import tvm.script @@ -1019,6 +1020,16 @@ def test( assert_structural_equal(test, from_source(code)) +def test_buffer_shape_repeated_var_prints_out_of_line(): + n = tvm.tirx.Var("n", "int32") + buffer = tvm.tirx.decl_buffer((n + n,), name="A") + func = tvm.tirx.PrimFunc([buffer], tvm.tirx.Evaluate(0)) + + code = func.script() + assert "n = T.int32()" in code + assert_structural_equal(func, from_source(code)) + + def test_kwargs_op_call(): # fmt: off @T.prim_func(private=True) @@ -1254,7 +1265,7 @@ def from_tuple(x: T.int32, y: T.float32) -> T.int32: def tuple_value(func): visited = [] - tvm.tirx.stmt_functor.post_order_visit(func.body, visited.append) + tvm_ffi.structural_walk(func.body, visited.append) bind = next(node for node in visited if isinstance(node, tvm.tirx.Bind)) return bind.value @@ -1366,7 +1377,7 @@ def _visit(node): if isinstance(node, tvm.tirx.DeclBuffer | tvm.tirx.AllocBuffer): bufs[node.buffer.name] = node.buffer - tvm.tirx.stmt_functor.post_order_visit(func.body, _visit) + tvm_ffi.structural_walk(func.body, _visit) return bufs @@ -1378,7 +1389,7 @@ def _visit(node): if isinstance(node, tvm.tirx.DeclBuffer): sources[node.buffer.name] = node.data - tvm.tirx.stmt_functor.post_order_visit(func.body, _visit) + tvm_ffi.structural_walk(func.body, _visit) return sources @@ -1698,7 +1709,7 @@ def func() -> None: # fmt: on binds = [] - tvm.tirx.stmt_functor.post_order_visit( + tvm_ffi.structural_walk( func.body, lambda node: binds.append(node) if isinstance(node, tvm.tirx.Bind) else None ) assert len(binds) == 1 @@ -1735,7 +1746,7 @@ def func() -> None: func = tvm.script.from_source(source, extra_vars={"T": T, "ptr": object()}) binds = [] - tvm.tirx.stmt_functor.post_order_visit( + tvm_ffi.structural_walk( func.body, lambda node: binds.append(node) if isinstance(node, tvm.tirx.Bind) else None ) assert len(binds) == 1 @@ -2731,7 +2742,7 @@ def func(): assert from_source(code).script() == code assert_structural_equal(func, from_source(code)) decls = [] - tvm.tirx.stmt_functor.post_order_visit( + tvm_ffi.structural_walk( func.body, lambda node: decls.append(node) if isinstance(node, tvm.tirx.DeclBuffer) else None, ) @@ -2990,7 +3001,7 @@ def func(A_ptr: T.handle): # fmt: on scope_defs = [] - tvm.tirx.stmt_functor.post_order_visit( + tvm_ffi.structural_walk( func.body, lambda s: ( scope_defs.append(getattr(s, "def")) if isinstance(s, tvm.tirx.ScopeIdDefStmt) else None @@ -3059,7 +3070,7 @@ def func(A_ptr: T.handle): # fmt: on scope_defs = [] - tvm.tirx.stmt_functor.post_order_visit( + tvm_ffi.structural_walk( func.body, lambda s: ( scope_defs.append(getattr(s, "def")) if isinstance(s, tvm.tirx.ScopeIdDefStmt) else None diff --git a/tests/python/tirx/transform/test_stmt_functor.py b/tests/python/tirx/transform/test_stmt_functor.py index 92ed2c47653d..8722473c57fe 100644 --- a/tests/python/tirx/transform/test_stmt_functor.py +++ b/tests/python/tirx/transform/test_stmt_functor.py @@ -1228,7 +1228,9 @@ def visit_buffer_load_(self, op): def test_op_call_nested_config_visited_and_substituted(): """Nested selector arrays participate in the core visitor and mutator.""" - from tvm.tirx.stmt_functor import post_order_visit, substitute + from tvm_ffi import structural_walk as post_order_visit + + from tvm.tirx.stmt_functor import substitute @T.prim_func def selector( diff --git a/tests/python/tirx/transform/test_transform_flatten_buffer.py b/tests/python/tirx/transform/test_transform_flatten_buffer.py index b47030eddabe..3cec134d8837 100644 --- a/tests/python/tirx/transform/test_transform_flatten_buffer.py +++ b/tests/python/tirx/transform/test_transform_flatten_buffer.py @@ -23,6 +23,8 @@ and hoists dead variables into the kernel ABI. """ +import tvm_ffi + import tvm import tvm.testing from tvm.script import tirx as T @@ -36,7 +38,7 @@ def visit(node): if isinstance(node, tvm.tirx.AllocBuffer | tvm.tirx.DeclBuffer): defined.add(node.buffer) - tvm.tirx.stmt_functor.post_order_visit(func.body, visit) + tvm_ffi.structural_walk(func.body, visit) return defined @@ -56,7 +58,7 @@ def visit(node): if isinstance(node, tvm.ir.TensorLoad) and not is_defined(node.source): stale.append(f"{where}: load of {node.source.name}") - tvm.tirx.stmt_functor.post_order_visit(expr, visit) + tvm_ffi.structural_walk(expr, visit) def visit(node): if isinstance(node, tvm.ir.TensorLoad | tvm.tirx.BufferStore): @@ -71,7 +73,7 @@ def visit(node): if node.buffer.elem_offset is not None: check_expr(node.buffer.elem_offset, f"elem_offset of {node.buffer.name}") - tvm.tirx.stmt_functor.post_order_visit(func.body, visit) + tvm_ffi.structural_walk(func.body, visit) assert not stale, f"stale buffer references after FlattenBuffer: {stale}" @@ -122,9 +124,9 @@ def inner(sub): if isinstance(sub, tvm.ir.TensorLoad): found.append(sub) - tvm.tirx.stmt_functor.post_order_visit(node.indices[0], inner) + tvm_ffi.structural_walk(node.indices[0], inner) - tvm.tirx.stmt_functor.post_order_visit(after.body, visit) + tvm_ffi.structural_walk(after.body, visit) assert found, "expected the folded elem_offset load in the mbar store index" @@ -143,7 +145,7 @@ def collect_before(node): if isinstance(node, tvm.tirx.AllocBuffer): before_allocs[node.buffer.name] = node.buffer - tvm.tirx.stmt_functor.post_order_visit(before.body, collect_before) + tvm_ffi.structural_walk(before.body, collect_before) after = _flatten(before) preserved = [] @@ -152,7 +154,7 @@ def visit(node): if isinstance(node, tvm.tirx.AllocBuffer) and node.buffer.name in before_allocs: preserved.append(node.buffer.same_as(before_allocs[node.buffer.name])) - tvm.tirx.stmt_functor.post_order_visit(after.body, visit) + tvm_ffi.structural_walk(after.body, visit) assert preserved and all(preserved), "already-flat buffer identity was not preserved" diff --git a/tests/python/tirx/transform/test_transform_lower_tirx.py b/tests/python/tirx/transform/test_transform_lower_tirx.py index d9a571942de5..f65050f58721 100644 --- a/tests/python/tirx/transform/test_transform_lower_tirx.py +++ b/tests/python/tirx/transform/test_transform_lower_tirx.py @@ -16,6 +16,7 @@ # under the License. import pytest +import tvm_ffi import tvm import tvm.testing @@ -55,7 +56,7 @@ def collect(node): if isinstance(node, tvm.tirx.AttrStmt) and node.attr_key == "thread_extent": extents[str(node.node.thread_tag)] = int(node.value) - tvm.tirx.stmt_functor.post_order_visit(func.body, collect) + tvm_ffi.structural_walk(func.body, collect) return extents diff --git a/tests/python/tvmscript/test_tvmscript_parser_source.py b/tests/python/tvmscript/test_tvmscript_parser_source.py index 31c7581d9cb6..824131793227 100644 --- a/tests/python/tvmscript/test_tvmscript_parser_source.py +++ b/tests/python/tvmscript/test_tvmscript_parser_source.py @@ -21,6 +21,7 @@ import pytest import tvm_ffi +from tvm_ffi import structural_walk as post_order_visit import tvm import tvm.testing @@ -30,7 +31,6 @@ from tvm.script.parser.core.diagnostics import Source from tvm.script.tirx import tile as Tx from tvm.tirx.stmt import TilePrimitiveCall -from tvm.tirx.stmt_functor import post_order_visit def _tirx_source(func):