This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 2ef0ae2609 [REFACTOR][TIR] Use StructuralWalk in place of
PostOrderVisit (#20308)
2ef0ae2609 is described below
commit 2ef0ae260903a82026abf1fc0d646e99ea72a0db
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Sep 10 12:57:36 2026 -0400
[REFACTOR][TIR] Use StructuralWalk in place of PostOrderVisit (#20308)
This PR replaces TIR PostOrderVisit and PreOrderVisit usage with
StructuralWalk.
A tree walk now uses a typed callback, and pre-order boolean descent
becomes WalkResult::Advance or WalkResult::Skip. Sites retain identity
sets only where a consumer depends on once-per-object behavior; buffer
printer use counting instead uses def_region_kind to distinguish
definitions from uses.
tirx::PostOrderVisit, tirx::PreOrderVisit, their helper visitors, FFI
globals, and Python wrappers are removed. Python callers use
tvm_ffi.structural_walk.
---
include/tvm/tirx/stmt_functor.h | 19 ---
python/tvm/relax/utils.py | 2 +-
.../tvm/s_tir/dlight/analysis/common_analysis.py | 4 +-
python/tvm/s_tir/dlight/gpu/fallback.py | 4 +-
python/tvm/s_tir/dlight/gpu/general_reduction.py | 6 +-
python/tvm/tirx/stmt_functor.py | 30 ----
src/backend/trn/codegen/codegen_trn.cc | 49 +++---
src/relax/analysis/tir_op_pattern_kind.cc | 16 +-
src/relax/script/printer/dependent_type.cc | 11 +-
src/relax/transform/rewrite_cuda_graph.cc | 26 +--
.../analysis/sblock_buffer_access_lca_detector.cc | 28 +--
.../feature_extractor/per_store_feature.cc | 67 ++++----
.../postproc/rewrite_cooperative_fetch.cc | 39 +++--
.../meta_schedule/postproc/rewrite_tensorize.cc | 49 +++---
src/s_tir/schedule/analysis/analysis.cc | 187 ++++++++++-----------
src/s_tir/schedule/analysis/reducer.cc | 46 +++--
src/s_tir/schedule/primitive/cache_index.cc | 56 +++---
src/s_tir/schedule/primitive/cache_read_write.cc | 12 +-
src/s_tir/schedule/primitive/for_kind.cc | 21 ++-
.../schedule/primitive/layout_transformation.cc | 14 +-
src/s_tir/schedule/primitive/pad_einsum.cc | 12 +-
src/s_tir/schedule/primitive/read_write_at.cc | 13 +-
src/s_tir/schedule/trace.cc | 11 +-
.../transform/lower_cross_thread_reduction.cc | 87 +++++-----
.../transform/memhammer_intermediate_stage.cc | 34 ++--
.../transform/memhammer_tensorcore_rewrite.cc | 25 +--
src/s_tir/transform/renew_defs.cc | 15 +-
src/te/operation/compute_op.cc | 21 ++-
src/te/tensor.cc | 28 ++-
src/tirx/ir/index_map.cc | 35 ++--
src/tirx/ir/stmt_functor.cc | 104 +-----------
src/tirx/ir/tir_visitor_with_path.h | 13 +-
src/tirx/script/printer/buffer.cc | 48 +++---
src/tirx/script/printer/function.cc | 27 ++-
src/tirx/script/printer/utils.h | 28 +--
src/tirx/transform/ir_utils.cc | 18 +-
tests/cpp/ir_functor_test.cc | 92 +++++-----
.../relax/test_backend_dispatch_sort_scan.py | 3 +-
tests/python/relax/test_pipeline.py | 2 +-
.../s_tir/base/test_sblock_dependence_info.py | 4 +-
.../s_tir/schedule/test_tir_schedule_analysis.py | 6 +-
.../schedule/test_tir_schedule_block_scope.py | 2 +-
.../test_tir_schedule_state_cached_flags.py | 2 +-
.../schedule/test_tir_schedule_transform_layout.py | 25 +++
.../transform/test_s_tir_transform_hoist_if.py | 4 +-
.../test_s_tir_transform_inject_double_buffer.py | 8 +-
.../test_s_tir_transform_inject_ptx_async_copy.py | 2 +-
.../test_s_tir_transform_inject_ptx_ldg32.py | 6 +-
.../test_s_tir_transform_inject_virtual_thread.py | 12 +-
.../test_s_tir_transform_loop_partition.py | 3 +-
...test_s_tir_transform_lower_thread_all_reduce.py | 4 +-
...st_s_tir_transform_memhammer_lower_auto_copy.py | 3 +-
tests/python/tirx-base/test_tir_base.py | 2 +-
.../test_tir_transform_bf16_legalize.py | 13 +-
.../test_tir_transform_lower_intrin.py | 3 +-
.../test_tir_transform_make_packed_api.py | 7 +-
...est_tir_transform_pointer_value_type_rewrite.py | 4 +-
.../test_tir_transform_split_host_device.py | 3 +-
.../test_tir_transform_storage_rewrite.py | 21 +--
.../tirx-transform/test_tir_transform_vectorize.py | 3 +-
.../python/tirx/codegen/test_codegen_blackwell.py | 5 +-
tests/python/tirx/codegen/test_codegen_dsmem.py | 4 +-
tests/python/tirx/codegen/test_ptx_addr.py | 3 +-
tests/python/tirx/test_hint.py | 12 +-
tests/python/tirx/test_op_namespace_cleanup.py | 5 +-
tests/python/tirx/test_parser_printer.py | 27 ++-
tests/python/tirx/transform/test_stmt_functor.py | 4 +-
.../transform/test_transform_flatten_buffer.py | 16 +-
.../tirx/transform/test_transform_lower_tirx.py | 3 +-
.../tvmscript/test_tvmscript_parser_source.py | 2 +-
70 files changed, 766 insertions(+), 754 deletions(-)
diff --git a/include/tvm/tirx/stmt_functor.h b/include/tvm/tirx/stmt_functor.h
index 54cf329a0a..9aa466f603 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<void(const ffi::ObjectRef&)> fvisit);
-
/*!
* \brief Substitute the var specified by vmap.
* \param stmt The source statement to be substituted
@@ -546,16 +537,6 @@ TVM_DLL Stmt SubstituteWithDataTypeLegalization(
TVM_DLL PrimExpr SubstituteWithDataTypeLegalization(
PrimExpr expr, std::function<ffi::Optional<PrimExpr>(const Var&)> vmap);
-/*!
- * \brief Recursively visit a statement or expression in pre DFS order,
applying fvisit.
- * If fvisit returns false, it won't visit the children of the node.
- * \param stmt_or_expr The statement or expression to be visited.
- * \param fvisit The visitor function to be applied. If fvisit returns false,
it won't visit the
- * children of the node
- */
-TVM_DLL void PreOrderVisit(const ffi::ObjectRef& stmt_or_expr,
- const std::function<bool(const ffi::ObjectRef&)>&
fvisit);
-
/*!
* \brief Check if the statement contains the specified node type.
*
diff --git a/python/tvm/relax/utils.py b/python/tvm/relax/utils.py
index 78242bd070..2d9e42523f 100644
--- a/python/tvm/relax/utils.py
+++ b/python/tvm/relax/utils.py
@@ -183,7 +183,7 @@ def gen_call_tir_inputs(
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 9c29d49d79..1a697ef7ef 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 @@ import logging
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_vars_used_in_prim_expr(expr: tirx.Expr) ->
set[tirx.Var]:
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 0acb39a6ba..0c65566632 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 _has_internal_thread_env(stmt: tirx.Stmt) -> bool:
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 d3d758afc9..b79d644ac1 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 @@ class GeneralReduction(GPUScheduleRule):
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 eec1c184f9..439e04930b 100644
--- a/python/tvm/tirx/stmt_functor.py
+++ b/python/tvm/tirx/stmt_functor.py
@@ -980,36 +980,6 @@ class StmtExprMutator(StmtMutator, ExprMutator):
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 da0f6f744b..5646054de5 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 <tvm/ffi/extra/structural_visit.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/runtime/logging.h>
#include <tvm/tirx/transform.h>
@@ -309,14 +310,18 @@ std::string CodeGenTrainium::PrintIndices(const
Array<PrimExpr>& indices) {
ctx_.buffer_index = 0;
ctx_.used_var_cnt = 0;
for (size_t i = 0; i < indices.size(); ++i) {
- PreOrderVisit(indices[i], [&](const ffi::ObjectRef& node) {
- if (const auto* v = node.as<VarNode>()) {
- if (ctx_.tensorized_loop_vars.count(v)) {
- ctx_.used_var_cnt++;
- }
+ std::unordered_set<const ffi::Object*> visited;
+ auto walk_fn = [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+ const VarNode* v = var.get();
+ if (!visited.insert(v).second) {
+ return ffi::WalkResult::Advance();
}
- return true;
- });
+ if (ctx_.tensorized_loop_vars.count(v)) {
+ ctx_.used_var_cnt++;
+ }
+ return ffi::WalkResult::Advance();
+ };
+ ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(indices[i], walk_fn);
}
for (size_t i = 0; i < indices.size(); ++i) {
if (i != 0) {
@@ -515,22 +520,22 @@ void CodeGenTrainium::VisitExpr_(const CallNode* op,
std::ostream& os) { // NOL
LOG(FATAL) << "Trainium codegen does not support call to " << op->op;
}
if (ctx_.mask.defined()) {
- PreOrderVisit(ctx_.mask, [&](const ffi::ObjectRef& node) {
- if (const auto* v = node.as<VarNode>()) {
- if (ctx_.tensorized_loop_vars.count(v)) {
- TVM_FFI_ICHECK(ctx_.loopvar2dim.count(v))
- << "nki_dim must be specified for tensorized loop variables used
in mask. However, "
- "it is not specified for "
- << ffi::GetRef<Var>(v);
- auto dim_str = ctx_.loopvar2dim[v];
- TVM_FFI_ICHECK(dim_str == "P" || dim_str == "F")
- << "Only nki_dim = P or F is allowed for tensorized loop
variables used in mask. "
- "However, "
- << ffi::GetRef<Var>(v) << " has nki_dim = " << dim_str;
- }
+ auto walk_fn = [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+ const VarNode* v = var.get();
+ if (ctx_.tensorized_loop_vars.count(v)) {
+ TVM_FFI_ICHECK(ctx_.loopvar2dim.count(v))
+ << "nki_dim must be specified for tensorized loop variables used
in mask. However, "
+ "it is not specified for "
+ << ffi::GetRef<Var>(v);
+ auto dim_str = ctx_.loopvar2dim[v];
+ TVM_FFI_ICHECK(dim_str == "P" || dim_str == "F")
+ << "Only nki_dim = P or F is allowed for tensorized loop variables
used in mask. "
+ "However, "
+ << ffi::GetRef<Var>(v) << " has nki_dim = " << dim_str;
}
- return true;
- });
+ return ffi::WalkResult::Advance();
+ };
+ ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(ctx_.mask, walk_fn);
os << ", mask=" << PrintExpr(ctx_.mask);
}
os << ")";
diff --git a/src/relax/analysis/tir_op_pattern_kind.cc
b/src/relax/analysis/tir_op_pattern_kind.cc
index 5287e56d96..de8a8d06d9 100644
--- a/src/relax/analysis/tir_op_pattern_kind.cc
+++ b/src/relax/analysis/tir_op_pattern_kind.cc
@@ -19,6 +19,7 @@
#include <tvm/arith/iter_affine_map.h>
#include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/relax/analysis.h>
#include <tvm/relax/op_attr_types.h>
@@ -260,15 +261,14 @@ class PatternKindAnalyzer : public StmtExprVisitor {
return false;
}
}
+ auto walk_fn = [&](const tirx::Var& var) -> ffi::Expected<ffi::WalkResult>
{
+ if (auto prim_var = var.as<tirx::PrimVar>()) {
+ vars.erase(prim_var.value().get());
+ }
+ return ffi::WalkResult::Advance();
+ };
for (const PrimExpr& index : load->indices) {
- PreOrderVisit(index, [&](const ffi::ObjectRef& node) {
- if (auto var = node.as<tirx::PrimVar>()) {
- if (vars.count(var.value().get())) {
- vars.erase(var.value().get());
- }
- }
- return true;
- });
+ ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(index, walk_fn);
}
return !vars.empty();
}
diff --git a/src/relax/script/printer/dependent_type.cc
b/src/relax/script/printer/dependent_type.cc
index af1f126fce..606e37a63f 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 <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/tirx/stmt_functor.h>
#include "./utils.h"
@@ -44,13 +45,15 @@ ExprDoc PrintShapeVar(const PrimExpr& e, const AccessPath&
e_p, const IRDocsifie
// Step 2. Figure out if the PrimExpr contains at least a func var
bool func_var_mode = false;
if (f != nullptr) {
- tirx::PostOrderVisit(e, [f, &func_var_mode](const ffi::ObjectRef& obj) ->
void {
- if (auto var = obj.as<tirx::PrimVar>()) {
- if (f->func_vars->count(var.value().get())) {
+ auto walk_fn = [f, &func_var_mode](const tirx::Var& var) ->
ffi::Expected<ffi::WalkResult> {
+ if (auto prim_var = var.as<tirx::PrimVar>()) {
+ if (f->func_vars->count(prim_var.value().get())) {
func_var_mode = true;
}
}
- });
+ return ffi::WalkResult::Advance();
+ };
+ ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(e, walk_fn);
}
// Step 3. Stringify the PrimExpr if func var exists
bool is_bare_type_var = false;
diff --git a/src/relax/transform/rewrite_cuda_graph.cc
b/src/relax/transform/rewrite_cuda_graph.cc
index 8515d9df3e..99b0f2d2d6 100644
--- a/src/relax/transform/rewrite_cuda_graph.cc
+++ b/src/relax/transform/rewrite_cuda_graph.cc
@@ -50,6 +50,7 @@
* with `CUDAGraphRewriter`.
*/
#include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/relax/analysis.h>
@@ -483,17 +484,22 @@ class CUDAGraphRewritePlanner : public ExprVisitor {
[[maybe_unused]] std::vector<const VarNode*>* vars_collector =
nullptr,
std::vector<tirx::PrimVar>* tir_vars_collector = nullptr) {
bool is_static = true;
- tirx::PostOrderVisit(expr, [&](const ffi::ObjectRef& e) {
- if (auto var = e.as<tirx::PrimVar>()) {
- if (!capture_symbolic_vars_.count(var.value())) {
- is_static = false;
- return;
- }
- if (tir_vars_collector != nullptr) {
- tir_vars_collector->push_back(var.value());
- }
+ std::unordered_set<const ffi::Object*> visited;
+ auto walk_fn = [&](const tirx::Var& var) -> ffi::Expected<ffi::WalkResult>
{
+ auto prim_var = var.as<tirx::PrimVar>();
+ if (!prim_var || !visited.insert(prim_var.value().get()).second) {
+ return ffi::WalkResult::Advance();
+ }
+ if (!capture_symbolic_vars_.count(prim_var.value())) {
+ is_static = false;
+ return ffi::WalkResult::Advance();
}
- });
+ if (tir_vars_collector != nullptr) {
+ tir_vars_collector->push_back(prim_var.value());
+ }
+ return ffi::WalkResult::Advance();
+ };
+ ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(expr, walk_fn);
return is_static;
}
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 1a2fac9415..2ccacf0ca1 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 <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/stmt_functor.h>
@@ -150,12 +151,12 @@ class LCADetector : public StmtExprVisitor {
auto do_collect_itervar_scope = [this](const IterVar& itervar,
const PrimExpr& binding) -> const
ScopeInfo* {
const ScopeInfo* highest_scope = nullptr;
- PostOrderVisit(binding, [this, &highest_scope](const ffi::ObjectRef&
obj) {
- if (auto var = obj.as<PrimVar>()) {
- const VarNode* loop_var = var.value().get();
+ auto walk_fn = [this, &highest_scope](const Var& var) ->
ffi::Expected<ffi::WalkResult> {
+ if (auto prim_var = var.as<PrimVar>()) {
+ const VarNode* loop_var = prim_var.value().get();
auto it = loop_scope_map_.find(loop_var);
if (it == loop_scope_map_.end()) {
- return;
+ return ffi::WalkResult::Advance();
}
const ScopeInfo* scope = it->second->parent_scope_info;
if (highest_scope == nullptr) {
@@ -164,7 +165,9 @@ class LCADetector : public StmtExprVisitor {
highest_scope = scope;
}
}
- });
+ return ffi::WalkResult::Advance();
+ };
+ ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(binding, walk_fn);
return highest_scope;
};
@@ -200,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<PrimVar>()) {
- const VarNode* iter_var = var.value().get();
+ auto handle_itervar = [&opaque_var_scope,
+ &scope](const Var& var) ->
ffi::Expected<ffi::WalkResult> {
+ if (auto prim_var = var.as<PrimVar>()) {
+ const VarNode* iter_var = prim_var.value().get();
auto dom_scope_it = opaque_var_scope.find(iter_var);
if (dom_scope_it == opaque_var_scope.end()) {
- return;
+ return ffi::WalkResult::Advance();
}
// find the highest loop scope the accessed buffer index has
// loop carried dependencies to (via opaque iter var binding).
@@ -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<ffi::WalkOrder::kPostOrder>(range->min,
handle_itervar);
+ ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(range->min +
range->extent - 1,
+ handle_itervar);
}
// the scope should be above `highest_reduce_scope` for reduce output
buffer.
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 6a20f4aa4b..5fe771d5db 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 <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/transform.h>
@@ -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<TensorLoadNode>()) {
- found = true;
- }
- });
- return found;
+ auto walk_fn = [](const TensorLoad&) -> ffi::Expected<ffi::WalkResult> {
+ return ffi::WalkResult::Interrupt(ffi::VisitInterrupt(true));
+ };
+ auto result = ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(expr,
walk_fn);
+ return result.has_value() ? result.value()->value.cast<bool>() : 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<TensorLoadNode>()) {
- BufferVar buffer = load->source.as_or_throw<tvm::tirx::BufferVar>();
- 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<ffi::WalkResult> {
+ BufferVar buffer = load->source.as_or_throw<tvm::tirx::BufferVar>();
+ 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<ffi::WalkOrder::kPostOrder>(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<const VarNode*> region_vars;
+ auto walk_fn = [®ion_vars](const Var& var) ->
ffi::Expected<ffi::WalkResult> {
+ if (auto prim_var = var.as<PrimVar>()) {
+ 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<PrimVar>()) {
- region_vars.insert(var.value().get());
- }
- });
+ ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(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 ad5bb883b3..2cea200f41 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 <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/op.h>
#include <tvm/s_tir/stmt.h>
@@ -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<tirx::BufferStoreNode>()) {
- max_bytes = std::max(max_bytes, store->value.ty().StorageBytes());
- } else if (const auto* load = obj.as<TensorLoadNode>()) {
- max_bytes = std::max(max_bytes,
load->ty.as_or_throw<PrimType>().StorageBytes());
- } else if (const auto* call = obj.as<CallNode>()) {
- 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<size_t>(max_bytes, 8);
- }
- } else if (const auto* cast = obj.as<prim::CastNode>()) {
- max_bytes = std::max(max_bytes,
cast->ty.as_or_throw<PrimType>().StorageBytes());
+ auto visit_store = [&](const tirx::BufferStore& store) ->
ffi::Expected<ffi::WalkResult> {
+ max_bytes = std::max(max_bytes, store->value.ty().StorageBytes());
+ return ffi::WalkResult::Advance();
+ };
+ auto visit_load = [&](const TensorLoad& load) ->
ffi::Expected<ffi::WalkResult> {
+ max_bytes = std::max(max_bytes,
load->ty.as_or_throw<PrimType>().StorageBytes());
+ return ffi::WalkResult::Advance();
+ };
+ auto visit_call = [&](const Call& call) -> ffi::Expected<ffi::WalkResult> {
+ 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<size_t>(max_bytes, 8);
}
- });
+ return ffi::WalkResult::Advance();
+ };
+ auto visit_cast = [&](const prim::Cast& cast) ->
ffi::Expected<ffi::WalkResult> {
+ max_bytes = std::max(max_bytes,
cast->ty.as_or_throw<PrimType>().StorageBytes());
+ return ffi::WalkResult::Advance();
+ };
+ ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(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 c1b9409406..638706b0b4 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 <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/logging.h>
#include <tvm/s_tir/meta_schedule/postproc.h>
@@ -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<std::tuple<ffi::String, ffi::String,
std::function<void(s_tir::SBlockRV)>>>* jobs) {
- tirx::PostOrderVisit(func->body, [=, &jobs](const ffi::ObjectRef& obj) {
- if (const auto* block = obj.as<tirx::SBlockNode>()) {
- tirx::StmtSRef block_sref = sch->GetSRef(block);
- std::string block_name =
block_sref->StmtAs<tirx::SBlockNode>()->name_hint;
- if (ffi::Optional<ffi::String> intrin_name =
- s_tir::GetAnn<ffi::String>(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<SBlockRV> child_blocks = sch->GetChildBlocks(block);
- TVM_FFI_ICHECK(child_blocks.size() == 1);
- ffi::Array<LoopRV> 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<ffi::WalkResult> {
+ tirx::StmtSRef block_sref = sch->GetSRef(block.get());
+ std::string block_name = block_sref->StmtAs<tirx::SBlockNode>()->name_hint;
+ if (ffi::Optional<ffi::String> intrin_name =
+ s_tir::GetAnn<ffi::String>(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<SBlockRV> child_blocks = sch->GetChildBlocks(block);
+ TVM_FFI_ICHECK(child_blocks.size() == 1);
+ ffi::Array<LoopRV> 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<ffi::WalkOrder::kPostOrder>(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 229c03e1ca..342632b6ce 100644
--- a/src/s_tir/schedule/analysis/analysis.cc
+++ b/src/s_tir/schedule/analysis/analysis.cc
@@ -17,6 +17,7 @@
* under the License.
*/
#include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/op.h>
#include <tvm/s_tir/stmt.h>
@@ -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<PrimVar>()) {
- vars->insert(var.value().get());
+ auto walk_fn = [vars](const Var& var) ->
ffi::Expected<ffi::WalkResult> {
+ if (auto prim_var = var.as<PrimVar>()) {
+ vars->insert(prim_var.value().get());
}
- });
+ return ffi::WalkResult::Advance();
+ };
+ ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(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<SBlockRealizeNode>()) {
- 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<ffi::WalkResult> {
+ 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<ffi::WalkResult> {
+ if (var.same_as(loop_var)) {
+ (*ref) += 1;
+ }
+ return ffi::WalkResult::Advance();
+ };
+ ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(binding, walk_fn);
}
- return true;
+ return ffi::WalkResult::Skip();
};
- PreOrderVisit(loop->body, f_visit);
+ ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(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<Op>& 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<CallNode>()) {
- if (op_set.count(call->op.operator->())) {
- found = true;
- }
+ auto walk_fn = [&op_set](const Call& call) -> ffi::Expected<ffi::WalkResult>
{
+ 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<ffi::WalkOrder::kPreOrder>(stmt, walk_fn);
+ return result.has_value() ? result.value()->value.cast<bool>() : 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<ffi::WalkResult> {
+ if (!is_one(realize->predicate)) {
+ return ffi::WalkResult::Interrupt(ffi::VisitInterrupt(true));
}
- if (const auto* realize = obj.as<SBlockRealizeNode>()) {
- // Case 1: BlockRealize
- if (!is_one(realize->predicate)) {
- has_branch = true;
- }
- } else if (obj->IsInstance<IfThenElseNode>() ||
obj->IsInstance<SelectNode>()) {
- // Case 2: IfThenElse / Select
- has_branch = true;
- } else if (const auto* call = obj.as<CallNode>()) {
- // 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<ffi::WalkResult> {
+ return ffi::WalkResult::Interrupt(ffi::VisitInterrupt(true));
+ };
+ auto visit_select = [](const Select&) -> ffi::Expected<ffi::WalkResult> {
+ return ffi::WalkResult::Interrupt(ffi::VisitInterrupt(true));
+ };
+ auto visit_call = [](const Call& call) -> ffi::Expected<ffi::WalkResult> {
+ 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<ffi::WalkOrder::kPreOrder>(stmt,
visit_realize, visit_branch,
+ visit_select,
visit_call);
+ return result.has_value() ? result.value()->value.cast<bool>() : false;
}
std::tuple</*exists=*/bool,
@@ -1583,22 +1578,16 @@ bool NeedsMultiLevelTiling(const ScheduleState& self,
const StmtSRef& block_sref
}
bool IsSpatialPrimFunc(const PrimFunc& func) {
- bool result = true;
- PreOrderVisit(func->body, [&result](const ffi::ObjectRef& obj) {
- if (result == false) {
- return false;
- }
- if (const auto* block = obj.as<SBlockNode>()) {
- 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<ffi::WalkResult> {
+ 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<ffi::WalkOrder::kPreOrder>(func->body,
walk_fn);
+ return result.has_value() ? result.value()->value.cast<bool>() : true;
}
std::pair<int64_t, int64_t> GetCumulativeSpaceAndReductionLength(const
s_tir::ScheduleState& self,
@@ -1738,23 +1727,20 @@ TensorIntrinDescInfo
ExtractTensorIntrinDescInfo(arith::AnalyzerObj* analyzer,
const auto* desc_scope_realize = desc_func->body.as<SBlockRealizeNode>();
TVM_FFI_ICHECK(desc_scope_realize);
{
- auto f_visit = [&](const ffi::ObjectRef& obj) -> bool {
- // Extract the block
- if (const auto* block = obj.as<SBlockRealizeNode>()) {
- info.desc_block = block;
- return false;
- }
- // Extract the loops
- if (const auto* loop = obj.as<ForNode>()) {
- 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<ffi::WalkResult> {
+ info.desc_block = block.get();
+ return ffi::WalkResult::Advance();
+ };
+ auto visit_loop = [&](const For& loop) -> ffi::Expected<ffi::WalkResult> {
+ 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<ffi::WalkOrder::kPostOrder>(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<ffi::WalkResult> {
+ if (auto prim_var = var.as<PrimVar>()) {
+ 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<PrimVar>()) {
- update_mask(var.value().get(), &lhs_buffer_masks,
lhs_buffer_index.at(lhs_buffer));
- }
- return true;
- });
+ ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(index, walk_fn);
}
}
diff --git a/src/s_tir/schedule/analysis/reducer.cc
b/src/s_tir/schedule/analysis/reducer.cc
index ac15128f7c..477b9b38fb 100644
--- a/src/s_tir/schedule/analysis/reducer.cc
+++ b/src/s_tir/schedule/analysis/reducer.cc
@@ -17,6 +17,7 @@
* under the License.
*/
#include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/te/operation.h>
#include "../utils.h"
@@ -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<SBlockNode>();
- 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<ffi::WalkResult> {
+ 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<ffi::WalkResult> {
+ // 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<AllocBufferNode>()) {
- buffer_allocated.insert(alloc->buffer.get());
- }
- const auto* store = obj.as<BufferStoreNode>();
- if (!store) {
- return true;
- }
-
+ buffer_allocated.insert(alloc->buffer.get());
+ return ffi::WalkResult::Advance();
+ };
+ auto visit_store = [&](const BufferStore& store) ->
ffi::Expected<ffi::WalkResult> {
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<ffi::WalkOrder::kPreOrder>(block->body,
visit_block,
+ visit_alloc,
visit_store);
+ return result.has_value() ? result.value()->value.cast<bool>() : 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 4c04ac5940..4afce512d1 100644
--- a/src/s_tir/schedule/primitive/cache_index.cc
+++ b/src/s_tir/schedule/primitive/cache_index.cc
@@ -18,6 +18,7 @@
*/
#include <tvm/arith/int_set.h>
#include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
#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>()) {
- 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<ffi::WalkResult> {
+ 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<ffi::WalkOrder::kPostOrder>(store->value, walk_fn);
auto cmp = [&](const PrimExpr& lhs, const PrimExpr& rhs) -> bool {
return CalculateExprComplexity(lhs) > CalculateExprComplexity(rhs);
};
@@ -235,9 +236,10 @@ ffi::Array<SBlock> 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<PrimVar>()) {
- Var iter_var = var.value();
+ auto collect_origin_var = [&info,
+ &expr_index](const Var& var) ->
ffi::Expected<ffi::WalkResult> {
+ if (auto prim_var = var.as<PrimVar>()) {
+ Var iter_var = prim_var.value();
const ffi::Array<Var>& 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<SBlock> MakeIndexCacheStage(IndexInfo* info,
const ffi::String& stora
info->origin_block_vars[expr_index].push_back(iter_var);
}
}
- });
+ return ffi::WalkResult::Advance();
+ };
+ ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(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<Var> 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<PrimVar>()) {
- 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<ffi::WalkResult> {
+ if (auto prim_var = var.as<PrimVar>()) {
+ 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<ffi::WalkOrder::kPostOrder>(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 d67d2486d5..878941a4bb 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 <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include <unordered_set>
@@ -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<Var> covered;
+ auto walk_fn = [&covered](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+ 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<Var>()) {
- covered.insert(var.value());
- }
- return true;
- });
+ ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(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 c2c7c1c887..833877af21 100644
--- a/src/s_tir/schedule/primitive/for_kind.cc
+++ b/src/s_tir/schedule/primitive/for_kind.cc
@@ -17,6 +17,7 @@
* under the License.
*/
#include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include "../utils.h"
@@ -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<SBlockRealizeNode>()) {
- // 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<SBlockRealize>(realize),
thread_scope);
+ auto walk_fn = [&](const SBlockRealize& realize) ->
ffi::Expected<ffi::WalkResult> {
+ // 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<ffi::WalkOrder::kPreOrder>(loop, walk_fn);
}
/*!
diff --git a/src/s_tir/schedule/primitive/layout_transformation.cc
b/src/s_tir/schedule/primitive/layout_transformation.cc
index 696821f322..4fc6ea3dc8 100644
--- a/src/s_tir/schedule/primitive/layout_transformation.cc
+++ b/src/s_tir/schedule/primitive/layout_transformation.cc
@@ -19,6 +19,7 @@
#include <tvm/arith/analyzer.h>
#include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/runtime/logging.h>
#include <optional>
@@ -1273,21 +1274,22 @@ IterVarType DetectNewBlockIterType(
const std::unordered_map<const VarNode*, IterVarType>&
block_iter_type_map) {
IterVarType result{kOpaque};
bool found = false;
- PostOrderVisit(expr, [&](const ffi::ObjectRef& obj) {
- if (auto var = obj.as<PrimVar>()) {
- auto it = block_iter_type_map.find(var.value().get());
+ auto walk_fn = [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+ if (auto prim_var = var.as<PrimVar>()) {
+ 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<ffi::WalkOrder::kPostOrder>(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 e07d1c550c..6f08e80fb4 100644
--- a/src/s_tir/schedule/primitive/pad_einsum.cc
+++ b/src/s_tir/schedule/primitive/pad_einsum.cc
@@ -18,6 +18,7 @@
*/
#include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/tirx/op.h>
#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<int>(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<ffi::WalkResult> {
if (node.get() == block) {
- found = true;
+ return ffi::WalkResult::Interrupt(ffi::VisitInterrupt(true));
}
- });
- if (found) {
+ return ffi::WalkResult::Advance();
+ };
+ auto result =
ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(scope_body[i], walk_fn);
+ if (result.has_value() && result.value()->value.cast<bool>()) {
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 b7352c272a..ef40ba230f 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 <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/s_tir/stmt.h>
#include <string>
@@ -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<SBlockRealizeNode>();
- if (realize == nullptr) {
- return true;
- }
+ &scope](const SBlockRealize& realize) ->
ffi::Expected<ffi::WalkResult> {
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<StmtSRef>(self_->stmt2ref.at(block)->parent),
/*high_exclusive=*/loop_sref_,
/*extra_relax_scope=*/scope)),
- /*bindings=*/GetBindings(ffi::GetRef<SBlockRealize>(realize)),
+ /*bindings=*/GetBindings(realize),
/*relaxed_regions=*/&relaxed_regions);
}
- return false;
+ return ffi::WalkResult::Skip();
};
- PreOrderVisit(subtrees[i], f_visit);
+ ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(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 a8272f0c94..acaeba7aca 100644
--- a/src/s_tir/schedule/trace.cc
+++ b/src/s_tir/schedule/trace.cc
@@ -17,6 +17,7 @@
* under the License.
*/
#include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/reflection/registry.h>
#include <sstream>
@@ -583,11 +584,11 @@ Trace TraceNode::Simplified(bool remove_postproc) const {
used_rvs.insert(obj.as<ffi::Object>());
continue;
} else if (auto prim_expr = obj.as<PrimExpr>()) {
- PostOrderVisit(*prim_expr, [&used_rvs](const ffi::ObjectRef& obj) ->
void {
- if (obj.as<VarNode>()) {
- used_rvs.insert(obj.get());
- }
- });
+ auto walk_fn = [&used_rvs](const Var& var) ->
ffi::Expected<ffi::WalkResult> {
+ used_rvs.insert(var.get());
+ return ffi::WalkResult::Advance();
+ };
+ ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(*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 474b23b2f9..47c87dfeb9 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 <tvm/arith/analyzer.h>
#include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/s_tir/stmt.h>
#include <tvm/s_tir/transform.h>
@@ -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<const VarNode*, int> buffer_writer_cnt;
- PreOrderVisit(scope_block->body, [&buffer_writer_cnt](const ffi::ObjectRef&
obj) {
- if (const auto* block = obj.as<SBlockNode>()) {
- 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<ffi::WalkResult> {
+ for (const BufferRegion& buffer_region : block->writes) {
+ ++buffer_writer_cnt[buffer_region->buffer.get()];
}
- return true;
- });
+ return ffi::WalkResult::Skip();
+ };
+ ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(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<AndNode>()) {
- ffi::Array<PrimExpr> sub_exprs = {and_node->a,
and_node->b};
- for (PrimExpr sub_expr : sub_exprs) {
- if (sub_expr->IsInstance<AndNode>()) {
- continue;
- }
- bool is_reduction = [sub_expr,
&reduction_loop_vars]() {
- ffi::Array<Var> 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<const ffi::Object*> visited_predicate_nodes;
+ auto walk_fn = [&wb_predicate, &reduction_loop_vars,
&visited_predicate_nodes](
+ const And& and_expr) -> ffi::Expected<ffi::WalkResult> {
+ if (!visited_predicate_nodes.insert(and_expr.get()).second) {
+ return ffi::WalkResult::Advance();
+ }
+ ffi::Array<PrimExpr> sub_exprs = {and_expr->a, and_expr->b};
+ for (PrimExpr sub_expr : sub_exprs) {
+ if (sub_expr->IsInstance<AndNode>()) {
+ continue;
+ }
+ bool is_reduction = [sub_expr, &reduction_loop_vars]() {
+ ffi::Array<Var> 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<ffi::WalkOrder::kPostOrder>(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<For>(reduction_loops[0]), [block, &visit](const
ffi::ObjectRef& obj) {
- if (const auto* realize = obj.as<SBlockRealizeNode>()) {
- 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<ffi::WalkResult> {
+ 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::WalkOrder::kPreOrder>(ffi::GetRef<For>(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 20272f87d2..3ce33fb6f8 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 <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include "memhammer_rewrite_rule.h"
@@ -282,26 +283,25 @@ std::pair<Stmt, SeqStmt> 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<TensorLoadNode>()) {
- if (buffer_load->source.as_or_throw<tvm::tirx::BufferVar>().scope() ==
"wmma.accumulator" ||
- buffer_load->source.as_or_throw<tvm::tirx::BufferVar>().scope() ==
"m16n8k8.matrixC") {
- if (target_buffer_load == nullptr) {
- target_buffer_load = buffer_load;
- } else {
-
TVM_FFI_ICHECK(target_buffer_load->source.as_or_throw<tvm::tirx::BufferVar>().same_as(
- buffer_load->source.as_or_throw<tvm::tirx::BufferVar>()))
- << "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<ffi::WalkResult> {
+ if (buffer_load->source.as_or_throw<tvm::tirx::BufferVar>().scope() ==
"wmma.accumulator" ||
+ buffer_load->source.as_or_throw<tvm::tirx::BufferVar>().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<tvm::tirx::BufferVar>().same_as(
+ buffer_load->source.as_or_throw<tvm::tirx::BufferVar>()))
+ << "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<ffi::WalkOrder::kPreOrder>(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 2de9369125..29092d821b 100644
--- a/src/s_tir/transform/memhammer_tensorcore_rewrite.cc
+++ b/src/s_tir/transform/memhammer_tensorcore_rewrite.cc
@@ -18,6 +18,7 @@
*/
#include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ir/op.h>
#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<TensorLoadNode>();
- if (load && load->source.as_or_throw<tvm::tirx::BufferVar>().scope() ==
"wmma.accumulator") {
+ auto walk_fn = [&](const TensorLoad& load) -> ffi::Expected<ffi::WalkResult>
{
+ if (load->source.as_or_throw<tvm::tirx::BufferVar>().scope() ==
"wmma.accumulator") {
TVM_FFI_ICHECK(buf_load == nullptr ||
buf_load->source.as_or_throw<tvm::tirx::BufferVar>().same_as(
load->source.as_or_throw<tvm::tirx::BufferVar>()))
<< "More than one source buffer of wmma accumulator found";
- buf_load = load;
+ buf_load = load.get();
}
- return true;
- });
+ return ffi::WalkResult::Advance();
+ };
+ ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(buf_store->value, walk_fn);
BufferVar src_buffer = buf_load->source.as_or_throw<tvm::tirx::BufferVar>();
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<TensorLoadNode>();
- if (load && load->source.as_or_throw<tvm::tirx::BufferVar>().scope() ==
"m16n8k8.matrixC") {
+ auto walk_fn = [&](const TensorLoad& load) -> ffi::Expected<ffi::WalkResult>
{
+ if (load->source.as_or_throw<tvm::tirx::BufferVar>().scope() ==
"m16n8k8.matrixC") {
TVM_FFI_ICHECK(buf_load == nullptr ||
buf_load->source.as_or_throw<tvm::tirx::BufferVar>().same_as(
load->source.as_or_throw<tvm::tirx::BufferVar>()))
<< "More than one source buffer of mma accumulator found";
- buf_load = load;
+ buf_load = load.get();
}
- return true;
- });
+ return ffi::WalkResult::Advance();
+ };
+ ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(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 212fd3c224..86fe153054 100644
--- a/src/s_tir/transform/renew_defs.cc
+++ b/src/s_tir/transform/renew_defs.cc
@@ -23,6 +23,7 @@
*/
#include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/stmt_functor.h>
@@ -59,14 +60,14 @@ class RenewDefMutator : public StmtExprMutator {
for (const auto& param : func->params) {
if (auto opt_buffer = param.as<BufferVar>()) {
const BufferVar& buffer = opt_buffer.value();
+ auto walk_fn = [&generator](const Var& var) ->
ffi::Expected<ffi::WalkResult> {
+ 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<Var>()) {
- if (generator.remap_.count(var.value()) == 0) {
- generator.ReDefineVar(var.value());
- }
- }
- });
+ ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(e, walk_fn);
}
}
}
diff --git a/src/te/operation/compute_op.cc b/src/te/operation/compute_op.cc
index 4bd50d15bd..1600541d81 100644
--- a/src/te/operation/compute_op.cc
+++ b/src/te/operation/compute_op.cc
@@ -23,6 +23,7 @@
*/
#include <tvm/arith/analyzer.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/builtin.h>
@@ -163,16 +164,18 @@ TVM_FFI_STATIC_INIT_BLOCK() {
ffi::Array<Tensor> ComputeOpNode::InputTensors() const {
ffi::Array<Tensor> ret;
std::unordered_set<Tensor> visited;
- auto visit = [&ret, &visited](const PrimExpr& e) {
- tirx::PostOrderVisit(e, [&ret, &visited](const ffi::ObjectRef& n) {
- if (auto call = n.as<Call>(); 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<ffi::WalkResult> {
+ 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<ffi::WalkOrder::kPostOrder>(e, walk_fn);
};
for (const PrimExpr& e : body) {
if (const auto* reduce = e.as<te::ReduceNode>()) {
diff --git a/src/te/tensor.cc b/src/te/tensor.cc
index a617d54bd4..15d644a9fe 100644
--- a/src/te/tensor.cc
+++ b/src/te/tensor.cc
@@ -21,6 +21,8 @@
* \file tensor.cc
*/
#include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_mutate.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/te/operation.h>
@@ -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<TensorNode>()
@@ -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<TensorNode>()
+ .attr(refl::type_attr::kStructuralVisit,
reinterpret_cast<void*>(&TensorVisit))
+ .attr(refl::type_attr::kStructuralMutate,
reinterpret_cast<void*>(&TensorMutate))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(&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 e4e473990d..1f53c6395c 100644
--- a/src/tirx/ir/index_map.cc
+++ b/src/tirx/ir/index_map.cc
@@ -25,6 +25,7 @@
#include <tvm/arith/int_set.h>
#include <tvm/arith/iter_affine_map.h>
#include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/unique_name_supply.h>
#include <tvm/tirx/index_map.h>
@@ -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<const ffi::Object*> visited;
+ auto walk_fn = [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+ auto prim_var = var.as<PrimVar>();
+ if (!prim_var) {
+ return ffi::WalkResult::Advance();
+ }
+ if (!visited.insert(prim_var.value().get()).second) {
+ return ffi::WalkResult::Advance();
+ }
+ if (ffi::Optional<ffi::String> 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<PrimVar>();
- if (!var) {
- return;
- }
- if (visited.count(obj.get())) {
- return;
- }
- visited.emplace(obj.get());
- if (ffi::Optional<ffi::String> 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<ffi::WalkOrder::kPostOrder>(expr, walk_fn);
});
}
diff --git a/src/tirx/ir/stmt_functor.cc b/src/tirx/ir/stmt_functor.cc
index 251f9d217e..45cd939e91 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<void(const ffi::ObjectRef&)> 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<void(const ffi::ObjectRef&)> f_;
- std::unordered_set<const ffi::Object*> visited_;
-};
-
-void PostOrderVisit(const ffi::ObjectRef& node, std::function<void(const
ffi::ObjectRef&)> fvisit) {
- if (node.as<StmtNode>()) {
- IRApplyVisit visitor(fvisit);
- visitor(node.as_or_throw<Stmt>());
- } else {
- IRApplyVisit visitor(fvisit);
- visitor(node.as_or_throw<Expr>());
- }
-}
-
+// Implementations of Substitute
class IRSubstitute : public StmtExprMutator {
public:
explicit IRSubstitute(std::function<ffi::Optional<Expr>(const Var&)> vmap) :
vmap_(vmap) {}
@@ -854,48 +818,6 @@ Expr Substitute(Expr expr,
std::function<ffi::Optional<Expr>(const Var&)> vmap)
return IRSubstitute(std::move(vmap))(std::move(expr));
}
-void PreOrderVisit(const ffi::ObjectRef& stmt_or_expr,
- const std::function<bool(const ffi::ObjectRef&)>& fvisit) {
- class PreOrderVisitor : public StmtExprVisitor {
- public:
- explicit PreOrderVisitor(const std::function<bool(const ffi::ObjectRef&)>&
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<bool(const ffi::ObjectRef&)>& f_;
- std::unordered_set<const ffi::Object*> visited_;
- };
-
- PreOrderVisitor visitor(fvisit);
- if (auto stmt = stmt_or_expr.as<Stmt>()) {
- visitor(stmt.value());
- } else if (auto expr = stmt_or_expr.as<Expr>()) {
- 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<ffi::Optional<Expr>(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<bool>(); });
- })
- .def("tirx.Substitute", [](ffi::ObjectRef node, ffi::Map<Var, Expr>
vmap) -> ffi::ObjectRef {
- if (node->IsInstance<StmtNode>()) {
- return Substitute(node.as_or_throw<Stmt>(), vmap);
- } else {
- return Substitute(node.as_or_throw<Expr>(), vmap);
- }
- });
+ refl::GlobalDef().def("tirx.Substitute",
+ [](ffi::ObjectRef node, ffi::Map<Var, Expr> vmap) ->
ffi::ObjectRef {
+ if (node->IsInstance<StmtNode>()) {
+ return Substitute(node.as_or_throw<Stmt>(), vmap);
+ } else {
+ return Substitute(node.as_or_throw<Expr>(), vmap);
+ }
+ });
}
} // namespace tirx
diff --git a/src/tirx/ir/tir_visitor_with_path.h
b/src/tirx/ir/tir_visitor_with_path.h
index b14ad713f3..13007a32d2 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 <tvm/ffi/extra/structural_visit.h>
#include <tvm/ir/module.h>
#include <tvm/ir/scope_stack.h>
#include <tvm/runtime/logging.h>
@@ -273,13 +274,13 @@ class TIRVisitorWithPath : protected
ExprFunctor<void(const Expr&, ffi::reflecti
auto shape_path = path->Attr("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<Var>()) {
- 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<ffi::WalkResult> {
+ if (auto var_def = WithDefIfUndefined(var, dim_path)) {
+ context.push_back(std::move(var_def).value());
}
- });
+ return ffi::WalkResult::Advance();
+ };
+ ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(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 fa5b0b6058..c3c456f390 100644
--- a/src/tirx/script/printer/buffer.cc
+++ b/src/tirx/script/printer/buffer.cc
@@ -40,22 +40,25 @@ ffi::Map<ffi::String, ExprDoc> BufferAttrs(
// Step 0. Set up statistics
std::unordered_map<const ffi::Object*, int> use_count;
- auto update_use_count = [&](const Expr& e) {
- tirx::PostOrderVisit(e, [&](const ffi::ObjectRef& n) {
- if (const VarNode* var = n.as<VarNode>()) {
- ++use_count[var];
+ std::unordered_set<const ffi::Object*> def_seen;
+ auto count_buffer_var = [&](const Var& var,
+ TVMFFIDefRegionKind kind) ->
ffi::Expected<ffi::WalkResult> {
+ 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<ffi::WalkOrder::kPostOrder>(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<ffi::WalkResult> {
+ ++use_count[var.get()];
+ return ffi::WalkResult::Advance();
+ };
+ ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(data.value(),
count_data_var);
}
auto is_new_var = [&](const Expr& e) { return e->IsInstance<VarNode>() &&
!d->IsVarDefined(e); };
auto add_out_of_line_var_def = [&](const Var& var, const AccessPath& var_p) {
@@ -90,16 +93,15 @@ ffi::Map<ffi::String, ExprDoc> BufferAttrs(
bool contains_new_var = false;
bool contains_compound_shape_var = false;
std::unordered_set<Var> vars_in_shape;
- tirx::PostOrderVisit(e, [&](const ffi::ObjectRef& obj) {
- if (const auto* var_node = obj.as<VarNode>()) {
- Var var = ffi::GetRef<Var>(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<ffi::WalkResult> {
+ 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<ffi::WalkOrder::kPostOrder>(e, walk_fn);
if (is_new_var(e)) {
add_out_of_line_var_def(e.as_or_throw<Var>(), e_p);
}
diff --git a/src/tirx/script/printer/function.cc
b/src/tirx/script/printer/function.cc
index 554dba11c5..a461cddfc0 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<tirx::Var> stringify_shape_vars;
std::unordered_set<tirx::Var> stringify_compound_shape_vars;
std::unordered_set<tirx::Var> shape_vars;
+ auto walk_fn = [&](const tirx::Var& shape_var) ->
ffi::Expected<ffi::WalkResult> {
+ 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::VarNode>()) {
- tirx::Var shape_var = ffi::GetRef<tirx::Var>(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<ffi::WalkOrder::kPostOrder>(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 4807a37e84..83ddf213c6 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 <tvm/ffi/extra/structural_equal.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/script/printer/ir_docsifier.h>
@@ -117,21 +118,22 @@ inline void AsDocBody(const tirx::Stmt& stmt, AccessPath
p, TIRFrameNode* f, con
if (const auto* seq_stmt = stmt.as<tirx::SeqStmtNode>()) {
ffi::Array<tirx::Stmt> 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<TensorLoadNode>()) {
- if
(load->source.as_or_throw<tvm::tirx::BufferVar>().same_as(buffer)) {
- found = true;
- }
- } else if (const auto* call = node.as<CallNode>()) {
- if (call->op.same_as(tirx::builtin::masked_load()) &&
!call->args.empty()) {
- if (auto var = call->args[0].as<Var>(); var &&
var.value().same_as(buffer.var())) {
- found = true;
- }
+ auto visit_load = [&](const TensorLoad& load) ->
ffi::Expected<ffi::WalkResult> {
+ if (load->source.as_or_throw<tvm::tirx::BufferVar>().same_as(buffer)) {
+ return ffi::WalkResult::Interrupt(ffi::VisitInterrupt(true));
+ }
+ return ffi::WalkResult::Advance();
+ };
+ auto visit_call = [&](const Call& call) ->
ffi::Expected<ffi::WalkResult> {
+ if (call->op.same_as(tirx::builtin::masked_load()) &&
!call->args.empty()) {
+ if (auto var = call->args[0].as<Var>(); var &&
var.value().same_as(buffer.var())) {
+ return ffi::WalkResult::Interrupt(ffi::VisitInterrupt(true));
}
}
- });
- return found;
+ return ffi::WalkResult::Advance();
+ };
+ auto result = ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(value,
visit_load, visit_call);
+ return result.has_value() ? result.value()->value.cast<bool>() : 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 e41223d819..6e7e44b780 100644
--- a/src/tirx/transform/ir_utils.cc
+++ b/src/tirx/transform/ir_utils.cc
@@ -26,6 +26,7 @@
#include <tvm/arith/analyzer.h>
#include <tvm/arith/int_solver.h>
#include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/scope_stack.h>
#include <tvm/s_tir/stmt.h>
@@ -129,10 +130,12 @@ class IRConvertSSA final : public StmtExprMutator {
defined_.insert(var_ptr);
}
};
+ auto walk_fn = [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+ 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<Var>()) check_var(var.value());
- });
+ ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(dim, walk_fn);
}
for (const auto& stride : buffer.value()->strides) {
if (auto var = stride.as<Var>()) check_var(var.value());
@@ -793,9 +796,10 @@ ffi::Optional<arith::IntConstraints>
ConditionalBoundsContext::TrySolveCondition
e->IsInstance<prim::EQNode>() || e->IsInstance<prim::NENode>()) {
bool is_simple = true;
std::vector<PrimVar> 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<ffi::WalkResult> {
if (obj.same_as(e)) {
- return;
+ return ffi::WalkResult::Advance();
} else if (const VarNode* var = obj.as<VarNode>()) {
PrimType var_ty = var->ty.as_or_throw<PrimType>();
if (var_ty.MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt)) {
@@ -806,7 +810,9 @@ ffi::Optional<arith::IntConstraints>
ConditionalBoundsContext::TrySolveCondition
obj->IsInstance<prim::MulNode>() ||
obj->IsInstance<prim::FloorDivNode>() ||
obj->IsInstance<prim::FloorModNode>() ||
obj->IsInstance<IntImmNode>();
}
- });
+ return ffi::WalkResult::Advance();
+ };
+ ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(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 c197b4931c..192f7f5abd 100644
--- a/tests/cpp/ir_functor_test.cc
+++ b/tests/cpp/ir_functor_test.cc
@@ -34,6 +34,7 @@
#include <tvm/tirx/stmt_functor.h>
#include <initializer_list>
+#include <unordered_set>
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<VarNode>()) ++n_var;
- });
+ std::unordered_set<const ffi::Object*> visited;
+ auto walk_fn = [&n_var, &visited](const Var& var) ->
ffi::Expected<ffi::WalkResult> {
+ if (visited.insert(var.get()).second) {
+ ++n_var;
+ }
+ return ffi::WalkResult::Advance();
+ };
+ ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(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<IfThenElseNode>()) {
- init_visited = true;
- return false;
- }
- if (const auto* eval = n.as<EvaluateNode>()) {
- if (const auto* int_imm = eval->value.as<IntImmNode>()) {
- 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<ffi::WalkResult> {
+ init_visited = true;
+ return ffi::WalkResult::Skip();
+ };
+ auto visit_evaluate = [&](const Evaluate& eval) ->
ffi::Expected<ffi::WalkResult> {
+ if (const auto* int_imm = eval->value.as<IntImmNode>()) {
+ 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<ffi::WalkOrder::kPreOrder>(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<ffi::WalkOrder::kPostOrder>(input,
expand_one).cast<Stmt>();
+ Stmt mapped =
+ ffi::StructuralMap<ffi::WalkOrder::kPostOrder>(input,
expand_one).as_or_throw<Stmt>();
EXPECT_FALSE(mapped.same_as(input));
EXPECT_EQ(shared.as<SeqStmtNode>()->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<SeqStmtNode>()->seq.GetArrayObj();
- Stmt mapped =
- ffi::StructuralMap<ffi::WalkOrder::kPostOrder>(std::move(input),
expand_one).cast<Stmt>();
+ Stmt mapped =
ffi::StructuralMap<ffi::WalkOrder::kPostOrder>(std::move(input), expand_one)
+ .as_or_throw<Stmt>();
EXPECT_EQ(mapped.get(), original);
EXPECT_NE(mapped.as<SeqStmtNode>()->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<ffi::WalkOrder::kPostOrder>(ordinary_input, transform)
- .template cast<Stmt>();
+ .template as_or_throw<Stmt>();
Stmt inplace_input = make_boundary_input();
const auto* original_root = inplace_input.get();
const auto* original_array =
inplace_input.as<SeqStmtNode>()->seq.GetArrayObj();
Stmt inplace =
ffi::StructuralMap<ffi::WalkOrder::kPostOrder>(std::move(inplace_input),
transform)
- .template cast<Stmt>();
+ .template as_or_throw<Stmt>();
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<ffi::WalkOrder::kPostOrder>(ordinary_input,
remove_all).cast<Stmt>();
+ Stmt ordinary =
ffi::StructuralMap<ffi::WalkOrder::kPostOrder>(ordinary_input, remove_all)
+ .as_or_throw<Stmt>();
Stmt inplace_input = make_boundary_input();
Stmt inplace =
ffi::StructuralMap<ffi::WalkOrder::kPostOrder>(std::move(inplace_input),
remove_all)
- .cast<Stmt>();
+ .as_or_throw<Stmt>();
EXPECT_TRUE(ffi::StructuralEqual()(ordinary, inplace));
for (const Stmt& result : {ordinary, inplace}) {
const auto* evaluate = result.as<EvaluateNode>();
@@ -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<ffi::WalkOrder::kPostOrder>(ordinary_input,
keep_last).cast<Stmt>();
+ ordinary =
+ ffi::StructuralMap<ffi::WalkOrder::kPostOrder>(ordinary_input,
keep_last).as_or_throw<Stmt>();
inplace_input = make_boundary_input();
inplace =
ffi::StructuralMap<ffi::WalkOrder::kPostOrder>(std::move(inplace_input),
keep_last)
- .cast<Stmt>();
+ .as_or_throw<Stmt>();
EXPECT_TRUE(ffi::StructuralEqual()(ordinary, inplace));
for (const Stmt& result : {ordinary, inplace}) {
const auto* evaluate = result.as<EvaluateNode>();
@@ -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<SeqStmtNode>()->seq[0].as<EvaluateNode>();
- Stmt mapped =
- ffi::StructuralMap<ffi::WalkOrder::kPostOrder>(std::move(input),
replace_one).cast<Stmt>();
+ Stmt mapped =
ffi::StructuralMap<ffi::WalkOrder::kPostOrder>(std::move(input), replace_one)
+ .as_or_throw<Stmt>();
const auto* mapped_seq = mapped.as<SeqStmtNode>();
ASSERT_NE(mapped_seq, nullptr);
@@ -537,8 +544,8 @@ TEST(IRF, StructuralMapPreservesSeqStmtElementUniqueness) {
ffi::Array<Stmt> shared_seq = {Evaluate(IntImm::Int32(1)),
Evaluate(IntImm::Int32(3))};
const auto* shared_first = shared_seq[0].as<EvaluateNode>();
Stmt input = SeqStmt(shared_seq);
- Stmt mapped =
- ffi::StructuralMap<ffi::WalkOrder::kPostOrder>(std::move(input),
replace_one).cast<Stmt>();
+ Stmt mapped =
ffi::StructuralMap<ffi::WalkOrder::kPostOrder>(std::move(input), replace_one)
+ .as_or_throw<Stmt>();
const auto* mapped_seq = mapped.as<SeqStmtNode>();
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<ffi::WalkOrder::kPostOrder>(std::move(unchanged_input),
no_float_match)
- .cast<Stmt>();
+ .as_or_throw<Stmt>();
EXPECT_EQ(unchanged.get(), unchanged_root);
EXPECT_TRUE(unchanged.as<SeqStmtNode>()->seq.same_as(shared_seq));
}
@@ -577,14 +584,13 @@ TEST(IRF, StructuralHooksPreserveScopeIdDefRegions) {
int binder = -1;
int extent = -1;
int preferred = -1;
- ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(
- make_input(),
- [&](const Var& var, TVMFFIDefRegionKind kind) ->
ffi::Expected<ffi::WalkResult> {
- 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<ffi::WalkResult> {
+ if (var->name == "binder") binder = kind;
+ if (var->name == "extent") extent = kind;
+ if (var->name == "preferred") preferred = kind;
+ return ffi::WalkResult::Advance();
+ };
+ ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(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 7fcd264917..df4aca033a 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 test_dispatch_cumsum_webgpu_axes_and_dtypes(
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 ae7a889940..dd8a15f819 100644
--- a/tests/python/relax/test_pipeline.py
+++ b/tests/python/relax/test_pipeline.py
@@ -176,7 +176,7 @@ def _has_thread_binding(func: tvm.tirx.PrimFunc) -> bool:
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 eb6ee6841d..6d20e16d49 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 gc
import sys
import pytest
+import tvm_ffi
import tvm
import tvm.testing
@@ -29,7 +30,6 @@ from tvm.s_tir import SBlockDependenceInfo
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 get_sblocks(func: PrimFunc):
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 48e7097b1b..3fb208e1af 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 import (
)
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 f98b45c4ec..f18e49cc48 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 02f56a156b..cccc4fd38d 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 4252f88f72..9cd1907cff 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 d59aedb4d2..8915aeebc8 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 verify_structure(stmt, expected_struct):
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 verify_structure(stmt, expected_struct):
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 5a761178b3..d8824ca6fc 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 test_double_buffer():
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 test_double_buffer():
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 test_double_buffer_transform():
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 a53f3c439c..922bcb992f 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 count_cp_async(stmt):
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 67e568a248..2c19027941 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 _count_alloc(stmt):
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 _count_ptx_ldg32(stmt):
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 471974f16c..315ec04de6 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 test_vthread():
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 test_vthread_extern():
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 test_vthread_if_then_else():
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 test_vthread_vectorized():
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 test_vthread_rewrites_masked_accesses():
}:
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 aa111bed1d..be1cfe5f90 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 @@ from tvm.script import tirx as T
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 5c08e66a95..17253785e1 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 _has_volatile_alloc_buffer(mod):
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 89a66cd4cc..5d050149f8 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_single_allocation(stmt, alloc_size=None):
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 dbf1eb623f..c82aad76bc 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 ce14378598..0a23e75b2e 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 test_bf16_masked_load_store_will_legalize():
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 052e2522f1..86d986250e 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 test_lower_nested_access_ptr():
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 81d46061ff..e48c5aac78 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 @@ Runtime error tests are in
tests/python/codegen/test_codegen_error_handling.py.
"""
import pytest
+import tvm_ffi
import tvm
import tvm.testing
@@ -37,7 +38,7 @@ def _find_compute_scope(func):
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 test_pointer_return():
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 test_buffer_alignment_attached_to_buffer_var():
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 89fd4011a6..e8a387cee3 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 test_decl_buffer_alias_chain_uses_flat_root_map():
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 2990efc725..ed5ce0f9a3 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 test_buffer_used_only_through_data_projection():
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 d5b271b468..b3133cba1e 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 test_alloc_seq():
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 test_alloc_different_dtypes():
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 test_address_of():
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 test_while_alloc():
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 test_alloc_seq_type():
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 test_alloc_seq_type2():
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 test_reuse_small_buffer():
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 test_decl_buffer_alias_extends_source_lifetime():
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 e7ac072cb7..e0940db94a 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 test_vectorize_nested_predicates_preserve_both_masks():
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 f02103a20b..26138ffd68 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 _assert_remote_mbarrier_ir(func, arrive_op_name,
n_arrives=1):
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
test_mbarrier_local_arrive_forwards_predicate_and_count():
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 d3ecd16a62..19481a5e3c 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 test_mapa_pointer_bind_codegen():
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 e333f9bb61..b26750370d 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 _calls(func, op_name):
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 1b3ca56487..6186c471e3 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 test_hint_statement():
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 test_hint_context_manager():
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 test_hint_with_attrs():
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 test_hint_no_message():
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 test_hint_access_buffer_region():
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 09ad39d4a6..78f56aa3d9 100644
--- a/tests/python/tirx/test_op_namespace_cleanup.py
+++ b/tests/python/tirx/test_op_namespace_cleanup.py
@@ -21,6 +21,7 @@ import sys
import types
import pytest
+import tvm_ffi
import tvm
from tvm.ir import Op, assert_structural_equal
@@ -36,7 +37,7 @@ def _tile_calls(func):
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 _expr_calls(func):
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 c6f6d5c21e..45ebe11b13 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_buffer():
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 test_tuple_let_binding_and_traversal():
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 _collect_buffers(func):
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 _collect_buffer_sources(func):
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 test_pointer_expression_assignment_uses_bind():
# 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 test_roundtrip_tmem_decl_buffer():
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 test_scope_id_dtype_uint32():
# 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 test_scope_id_dtype_uint32_deferred_extent():
# 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 92ed2c4765..8722473c57 100644
--- a/tests/python/tirx/transform/test_stmt_functor.py
+++ b/tests/python/tirx/transform/test_stmt_functor.py
@@ -1228,7 +1228,9 @@ def test_op_call_pointer_config_visited_and_mutated():
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 b47030edda..3cec134d88 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 @@ the rebuilt identity, or SplitHostDevice later sees them as
undefined
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 _collect_defined_buffers(func):
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 _assert_loads_reference_defined_buffers(func):
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 _assert_loads_reference_defined_buffers(func):
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 test_flatten_remaps_loads_in_folded_elem_offset():
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 test_flatten_keeps_identity_of_already_flat_buffers():
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 test_flatten_keeps_identity_of_already_flat_buffers():
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 d9a571942d..f65050f587 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 _launch_thread_extents(func):
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 31c7581d9c..8241317932 100644
--- a/tests/python/tvmscript/test_tvmscript_parser_source.py
+++ b/tests/python/tvmscript/test_tvmscript_parser_source.py
@@ -21,6 +21,7 @@ import inspect
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 import doc_core as doc
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):
