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 45713b39a3 [REFACTOR][TIRx][S-TIR] Localize SSA conversion and
simplify lexical remapping (#20695)
45713b39a3 is described below
commit 45713b39a3b1b87abe7ddf51fb1822fb6795f71c
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Oct 12 22:44:13 2026 -0400
[REFACTOR][TIRx][S-TIR] Localize SSA conversion and simplify lexical
remapping (#20695)
Keep SSA conversion local to TIRx and S-TIR, removing S-TIR's dependency
on the private TIRx converter while preserving its block-specific
traversal.
Replace remap histories and scope-exit callbacks with scope-local maps
of the latest variable binding. Definition tracking remains global
across scopes and functions; pattern definitions reuse only visible
bindings. Function attributes and metadata are preserved when parameters
or bodies change, and derived function signatures are refreshed.
---
include/tvm/ir/scope_stack.h | 4 +
src/s_tir/transform/ir_utils.cc | 234 ++++++++++++++++++++++++++++--
src/s_tir/transform/ir_utils.h | 3 +-
src/s_tir/transform/lower_opaque_block.cc | 1 +
src/tirx/transform/ir_utils.cc | 171 ++++++++++------------
src/tirx/transform/ir_utils.h | 85 -----------
6 files changed, 307 insertions(+), 191 deletions(-)
diff --git a/include/tvm/ir/scope_stack.h b/include/tvm/ir/scope_stack.h
index b5ea10656f..2fdc546ac4 100644
--- a/include/tvm/ir/scope_stack.h
+++ b/include/tvm/ir/scope_stack.h
@@ -57,6 +57,10 @@ class ScopeStack {
/*! \brief Construct with one initial scope level. */
ScopeStack() { stack_.emplace_back(); }
+ /*! \brief Iterate active scopes from innermost to outermost. */
+ auto rbegin() const { return stack_.rbegin(); }
+ auto rend() const { return stack_.rend(); }
+
/*! \brief Return the number of active scopes. */
size_t size() const { return stack_.size(); }
diff --git a/src/s_tir/transform/ir_utils.cc b/src/s_tir/transform/ir_utils.cc
index 326657dcb3..6b13342329 100644
--- a/src/s_tir/transform/ir_utils.cc
+++ b/src/s_tir/transform/ir_utils.cc
@@ -19,36 +19,39 @@
#include "ir_utils.h"
+#include <tvm/ffi/extra/structural_visit.h>
+#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/attrs.h>
#include <tvm/ir/prim/op.h>
+#include <tvm/ir/scope_stack.h>
#include <tvm/s_tir/stmt.h>
#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/sym/analyzer.h>
+#include <tvm/tirx/function.h>
#include <tvm/tirx/op/memory.h>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+
namespace tvm {
namespace s_tir {
using namespace tirx;
using namespace tvm::prim;
namespace {
-class SIRConvertSSA final : public tirx::IRConvertSSA {
+class SIRConvertSSA final : public StmtExprMutator {
public:
- TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(SIRConvertSSA,
tirx::IRConvertSSA)
- using tirx::IRConvertSSA::Mutate_;
-
- protected:
- static void InitVTable(VTable* table) {
- tirx::IRConvertSSA::InitVTable(table);
- SetDispatch<SIRConvertSSA, SBlockNode>(table);
- SetDispatch<SIRConvertSSA, SBlockRealizeNode>(table);
- }
+ TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(SIRConvertSSA, StmtExprMutator)
+ using StmtExprMutator::Mutate;
+ using StmtExprMutator::Mutate_;
+ Function VisitFunction(Function func);
+ IRModule VisitIRModule(IRModule mod);
- public:
UnchangedOr<Stmt> Mutate_(const SBlockNode* op, InplaceMode mode) {
SBlock block = ffi::GetRef<SBlock>(op);
- return WithScope([&]() -> Stmt {
+ return scope_.WithNewScope([&]() -> Stmt {
auto iter_vars = op->iter_vars.Map([&](IterVar iter) {
Var var = DefineVar(iter->var);
if (!var.same_as(iter->var)) iter.CopyOnWrite()->var =
var.as_or_throw<PrimVar>();
@@ -64,7 +67,214 @@ class SIRConvertSSA final : public tirx::IRConvertSSA {
UnchangedOr<Stmt> Mutate_(const SBlockRealizeNode* op, InplaceMode mode) {
return s_tir::StmtExprMutator::MutateBlockRealize(this, op, mode);
}
+
+ private:
+ UnchangedOr<Expr> Mutate_(const VarNode* op, InplaceMode inplace_mode) final;
+ Var DefineVar(Var var);
+ ffi::Optional<Var> GetVarRemap(const Var& var);
+ UnchangedOr<Stmt> Mutate_(const BindNode* op, InplaceMode inplace_mode)
final;
+ UnchangedOr<Stmt> Mutate_(const IfNode* op, InplaceMode inplace_mode) final;
+ UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final;
+ UnchangedOr<Stmt> Mutate_(const WhileNode* op, InplaceMode inplace_mode)
final;
+ UnchangedOr<Stmt> Mutate_(const RegionStmtNode* op, InplaceMode
inplace_mode) final;
+ UnchangedOr<PrimExpr> Mutate_(const prim::LetNode* op, InplaceMode
inplace_mode) final;
+ static Var MakeNewVar(const Var& old_var);
+ void SetVarRemap(const Var& old_var, const Var& new_var);
+
+ // Lexical visibility ends with each scope, but definitions stay unique
across the module.
+ ScopeStack<std::unordered_map<const VarNode*, Var>> scope_;
+ std::unordered_set<const VarNode*> defined_;
};
+
+Function SIRConvertSSA::VisitFunction(Function func) {
+ return scope_.WithNewScope([&]() -> Function {
+ std::unordered_set<const VarNode*> parameter_symbols;
+ // Define explicit parameters before the symbolic values in their types.
+ for (const Var& param : func->params) {
+ DefineVar(param);
+ parameter_symbols.insert(param.get());
+ }
+ for (const Var& param : func->params) {
+ ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(
+ param->ty, [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+ if (!parameter_symbols.count(var.get())) {
+ parameter_symbols.insert(var.get());
+ DefineVar(var);
+ }
+ return ffi::WalkResult::Advance();
+ });
+ }
+ auto params = func->params.Map([&](const Var& var) {
+ Var mapped = GetVarRemap(var).value_or(var);
+ Type type = Mutate(var->ty, InplaceMode::kDisallow)
+ .as_or_throw<UnchangedOr<Type>>()
+ .ValueOrUnchanged(var->ty);
+ if (!type.same_as(mapped->ty)) {
+ mapped = mapped.CopyWithType(type);
+ SetVarRemap(var, mapped);
+ }
+ return mapped;
+ });
+
+ auto body_result = Mutate(func->body, InplaceMode::kDisallow);
+ bool body_unchanged = body_result.UnchangedOrSameAs(func->body);
+ auto body = std::move(body_result).ValueOrUnchanged(func->body);
+
+ if (!params.same_as(func->params) || !body_unchanged) {
+ auto* node = func.CopyOnWrite();
+ node->params = std::move(params);
+ node->body = std::move(body);
+ node->ty = node->func_type_annotation();
+ }
+ return func;
+ });
+}
+
+UnchangedOr<Expr> SIRConvertSSA::Mutate_(const VarNode* op, InplaceMode
inplace_mode) {
+ Var var = ffi::GetRef<Var>(op);
+ if (def_region_kind() != kTVMFFIDefRegionKindNone) {
+ if (def_region_kind() == kTVMFFIDefRegionKindPattern) {
+ if (auto mapped = GetVarRemap(var)) return mapped.value();
+ }
+ return DefineVar(var);
+ }
+ Var mapped = GetVarRemap(var).value_or(var);
+ if (!mapped.same_as(var)) return mapped;
+ return ffi::Unchanged();
+}
+
+UnchangedOr<PrimExpr> SIRConvertSSA::Mutate_(const prim::LetNode* op,
InplaceMode inplace_mode) {
+ ffi::Any previous_remap = VarRemapGet(op->var);
+ PrimExpr result = scope_.WithNewScope([&] {
+ return StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<PrimExpr>(op));
+ });
+ VarRemapSet(op->var, previous_remap);
+ return result;
+}
+
+Var SIRConvertSSA::DefineVar(Var var) {
+ Type type = WithDefRegionKind(kTVMFFIDefRegionKindNone, [&] {
+ return Mutate(var->ty, InplaceMode::kDisallow)
+ .as_or_throw<UnchangedOr<Type>>()
+ .ValueOrUnchanged(var->ty);
+ });
+ Var result = defined_.count(var.get()) ? MakeNewVar(var) : var;
+ defined_.insert(var.get());
+ if (!type.same_as(result->ty)) result = result.CopyWithType(type);
+ SetVarRemap(var, result);
+ return result;
+}
+
+ffi::Optional<Var> SIRConvertSSA::GetVarRemap(const Var& var) {
+ for (auto it = scope_.rbegin(); it != scope_.rend(); ++it) {
+ if (auto found = it->find(var.get()); found != it->end()) return
found->second;
+ }
+ return std::nullopt;
+}
+
+UnchangedOr<Stmt> SIRConvertSSA::Mutate_(const BindNode* op, InplaceMode
inplace_mode) {
+ Var var = op->var;
+ ffi::Any previous_remap = VarRemapGet(var);
+ // The ordinary Bind path rewrites the RHS before defining the Var and
propagates its type.
+ Stmt stmt = StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
+ Var bound_var = stmt.as<BindNode>()->var;
+ if (!bound_var.same_as(GetVarRemap(var).value_or(var))) SetVarRemap(var,
bound_var);
+ VarRemapSet(var, previous_remap);
+ return stmt;
+}
+
+UnchangedOr<Stmt> SIRConvertSSA::Mutate_(const IfNode* op, InplaceMode
inplace_mode) {
+ // Each branch gets its own scope so Bind remaps in one branch
+ // do not leak into the other.
+ auto condition_result = Mutate(op->condition, inplace_mode);
+ bool condition_unchanged = condition_result.UnchangedOrSameAs(op->condition);
+ PrimExpr condition =
std::move(condition_result).ValueOrUnchanged(op->condition);
+ Stmt then_case = scope_.WithNewScope([&]() -> Stmt {
+ return Mutate(op->then_case, inplace_mode).ValueOrUnchanged(op->then_case);
+ });
+ ffi::Optional<SeqStmt> else_case;
+ if (op->else_case) {
+ else_case = scope_.WithNewScope([&]() -> SeqStmt {
+ return Mutate(op->else_case.value(),
inplace_mode).ValueOrUnchanged(op->else_case.value());
+ });
+ }
+ if (condition_unchanged && then_case.same_as(op->then_case) &&
else_case.same_as(op->else_case)) {
+ return ffi::Unchanged();
+ }
+ return If(condition, then_case, else_case);
+}
+
+UnchangedOr<Stmt> SIRConvertSSA::Mutate_(const ForNode* op, InplaceMode
inplace_mode) {
+ const Var& v = op->loop_var;
+ if (defined_.count(v.get())) {
+ return scope_.WithNewScope([&]() -> Stmt {
+ Var new_var = MakeNewVar(v);
+ SetVarRemap(v, new_var);
+ Stmt stmt =
+ StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
+ auto n = ffi::make_object<ForNode>(*stmt.as<ForNode>());
+ n->loop_var = new_var.as_or_throw<PrimVar>();
+ return For(n);
+ });
+ } else {
+ defined_.insert(v.get());
+ return scope_.WithNewScope([&]() -> Stmt {
+ SetVarRemap(v, v);
+ return StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
+ });
+ }
+}
+
+UnchangedOr<Stmt> SIRConvertSSA::Mutate_(const WhileNode* op, InplaceMode
inplace_mode) {
+ return scope_.WithNewScope([&]() -> Stmt {
+ return StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
+ });
+}
+
+UnchangedOr<Stmt> SIRConvertSSA::Mutate_(const RegionStmtNode* op, InplaceMode
inplace_mode) {
+ auto attrs = Mutate(op->attrs, InplaceMode::kDisallow)
+ .as_or_throw<UnchangedOr<DictAttrs>>()
+ .ValueOrUnchanged(op->attrs);
+ auto args = Mutate(op->args, InplaceMode::kDisallow)
+ .as_or_throw<UnchangedOr<ffi::Array<Expr>>>()
+ .ValueOrUnchanged(op->args);
+ ffi::Array<Var> params;
+ Stmt body = scope_.WithNewScope([&]() -> Stmt {
+ for (const Var& var : op->body_params) params.push_back(DefineVar(var));
+ return Mutate(op->body, InplaceMode::kDisallow).ValueOrUnchanged(op->body);
+ });
+ ffi::Array<Var> results;
+ for (const Var& var : op->result_vars) results.push_back(DefineVar(var));
+ return RegionStmt(op->op, std::move(args), std::move(params),
std::move(attrs), std::move(body),
+ std::move(results), op->loc);
+}
+
+Var SIRConvertSSA::MakeNewVar(const Var& old_var) {
+ return Var(old_var->name, old_var->ty, old_var->loc);
+}
+
+void SIRConvertSSA::SetVarRemap(const Var& old_var, const Var& new_var) {
+ scope_.Current().insert_or_assign(old_var.get(), new_var);
+}
+
+IRModule SIRConvertSSA::VisitIRModule(IRModule mod) {
+ ffi::Map<GlobalVar, BaseFunc> functions;
+ bool made_change = false;
+ for (auto [gvar, base_func] : mod->functions) {
+ if (auto* ptr = base_func.as<tirx::FunctionNode>()) {
+ auto updated = VisitFunction(ffi::GetRef<tirx::Function>(ptr));
+ if (!updated.same_as(base_func)) {
+ made_change = true;
+ base_func = updated;
+ }
+ }
+ functions.Set(gvar, base_func);
+ }
+ if (made_change) {
+ mod.CopyOnWrite()->functions = std::move(functions);
+ }
+ return mod;
+}
} // namespace
Stmt ConvertSSA(Stmt stmt) {
diff --git a/src/s_tir/transform/ir_utils.h b/src/s_tir/transform/ir_utils.h
index b23d30eb5b..9eddc15957 100644
--- a/src/s_tir/transform/ir_utils.h
+++ b/src/s_tir/transform/ir_utils.h
@@ -21,13 +21,12 @@
#define TVM_S_TIR_TRANSFORM_IR_UTILS_H_
#include <tvm/ffi/container/tuple.h>
+#include <tvm/ir/module.h>
#include <tvm/s_tir/stmt.h>
#include <string>
#include <unordered_map>
-#include "../../tirx/transform/ir_utils.h"
-
namespace tvm {
namespace s_tir {
diff --git a/src/s_tir/transform/lower_opaque_block.cc
b/src/s_tir/transform/lower_opaque_block.cc
index 1ae256fb15..a9e90650f3 100644
--- a/src/s_tir/transform/lower_opaque_block.cc
+++ b/src/s_tir/transform/lower_opaque_block.cc
@@ -23,6 +23,7 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/prim/op.h>
#include <tvm/ir/scope_stack.h>
#include <tvm/s_tir/stmt.h>
#include <tvm/s_tir/stmt_functor.h>
diff --git a/src/tirx/transform/ir_utils.cc b/src/tirx/transform/ir_utils.cc
index 536eb6da8e..6aa37aa957 100644
--- a/src/tirx/transform/ir_utils.cc
+++ b/src/tirx/transform/ir_utils.cc
@@ -43,83 +43,86 @@ namespace tvm {
namespace tirx {
using namespace tvm::prim;
-Function IRConvertSSA::VisitFunction(Function func) {
- std::unordered_set<const VarNode*> parameter_symbols;
- // Define explicit parameters before the symbolic values in their types.
- for (const Var& param : func->params) {
- DefineVar(param);
- parameter_symbols.insert(param.get());
- }
- for (const Var& param : func->params) {
- ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(
- param->ty, [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
- if (!parameter_symbols.count(var.get())) {
- parameter_symbols.insert(var.get());
- DefineVar(var);
- }
- return ffi::WalkResult::Advance();
- });
- }
- auto params = func->params.Map([&](const Var& var) {
- Var mapped = GetRemappedVar(var);
- Type type = Mutate(var->ty, InplaceMode::kDisallow)
- .as_or_throw<UnchangedOr<Type>>()
- .ValueOrUnchanged(var->ty);
- if (!type.same_as(mapped->ty)) {
- mapped = mapped.CopyWithType(type);
- PushVarRemap(var, mapped);
- }
- return mapped;
- });
+namespace {
+class IRConvertSSA final : public StmtExprMutator {
+ public:
+ TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(IRConvertSSA, StmtExprMutator)
+ using StmtExprMutator::Mutate;
+ using StmtExprMutator::Mutate_;
+ Function VisitFunction(Function func);
+ IRModule VisitIRModule(IRModule mod);
- auto attrs = [&]() -> DictAttrs {
- ffi::Map<ffi::String, ffi::Any> dict;
- bool made_change = false;
+ private:
+ UnchangedOr<Expr> Mutate_(const VarNode* op, InplaceMode inplace_mode) final;
+ Var DefineVar(Var var);
+ ffi::Optional<Var> GetVarRemap(const Var& var);
+ UnchangedOr<Stmt> Mutate_(const BindNode* op, InplaceMode inplace_mode)
final;
+ UnchangedOr<Stmt> Mutate_(const IfNode* op, InplaceMode inplace_mode) final;
+ UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final;
+ UnchangedOr<Stmt> Mutate_(const WhileNode* op, InplaceMode inplace_mode)
final;
+ UnchangedOr<Stmt> Mutate_(const RegionStmtNode* op, InplaceMode
inplace_mode) final;
+ UnchangedOr<PrimExpr> Mutate_(const prim::LetNode* op, InplaceMode
inplace_mode) final;
+ static Var MakeNewVar(const Var& old_var);
+ void SetVarRemap(const Var& old_var, const Var& new_var);
- for (const auto& [key, old_value] : func->attrs->dict) {
- auto value = old_value;
- if (auto expr = value.as<PrimExpr>()) {
- value = Mutate(expr.value(),
InplaceMode::kDisallow).ValueOrUnchanged(expr.value());
- } else if (auto* stmt = value.as<StmtNode>()) {
- value = Mutate(ffi::GetRef<Stmt>(stmt), InplaceMode::kDisallow)
- .ValueOrUnchanged(ffi::GetRef<Stmt>(stmt));
- }
+ // Lexical visibility ends with each scope, but definitions stay unique
across the module.
+ ScopeStack<std::unordered_map<const VarNode*, Var>> scope_;
+ std::unordered_set<const VarNode*> defined_;
+};
- made_change = made_change || !value.same_as(old_value);
- dict.Set(key, value);
+Function IRConvertSSA::VisitFunction(Function func) {
+ return scope_.WithNewScope([&]() -> Function {
+ std::unordered_set<const VarNode*> parameter_symbols;
+ // Define explicit parameters before the symbolic values in their types.
+ for (const Var& param : func->params) {
+ DefineVar(param);
+ parameter_symbols.insert(param.get());
}
-
- if (made_change) {
- return DictAttrs(dict);
- } else {
- return func->attrs;
+ for (const Var& param : func->params) {
+ ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(
+ param->ty, [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+ if (!parameter_symbols.count(var.get())) {
+ parameter_symbols.insert(var.get());
+ DefineVar(var);
+ }
+ return ffi::WalkResult::Advance();
+ });
}
- }();
+ auto params = func->params.Map([&](const Var& var) {
+ Var mapped = GetVarRemap(var).value_or(var);
+ Type type = Mutate(var->ty, InplaceMode::kDisallow)
+ .as_or_throw<UnchangedOr<Type>>()
+ .ValueOrUnchanged(var->ty);
+ if (!type.same_as(mapped->ty)) {
+ mapped = mapped.CopyWithType(type);
+ SetVarRemap(var, mapped);
+ }
+ return mapped;
+ });
- auto body_result = Mutate(func->body, InplaceMode::kDisallow);
- bool body_unchanged = body_result.UnchangedOrSameAs(func->body);
- auto body = std::move(body_result).ValueOrUnchanged(func->body);
+ auto body_result = Mutate(func->body, InplaceMode::kDisallow);
+ bool body_unchanged = body_result.UnchangedOrSameAs(func->body);
+ auto body = std::move(body_result).ValueOrUnchanged(func->body);
- // If anything changed, update the returned function
- if (!params.same_as(func->params) || !attrs.same_as(func->attrs) ||
!body_unchanged) {
- func = Function(params, body, func->ret_type, attrs);
- }
-
- // Pop function-scope remaps in reverse order
- PopAllRemapsInCurrentScope();
- return func;
+ if (!params.same_as(func->params) || !body_unchanged) {
+ auto* node = func.CopyOnWrite();
+ node->params = std::move(params);
+ node->body = std::move(body);
+ node->ty = node->func_type_annotation();
+ }
+ return func;
+ });
}
UnchangedOr<Expr> IRConvertSSA::Mutate_(const VarNode* op, InplaceMode
inplace_mode) {
Var var = ffi::GetRef<Var>(op);
if (def_region_kind() != kTVMFFIDefRegionKindNone) {
if (def_region_kind() == kTVMFFIDefRegionKindPattern) {
- Var mapped = GetRemappedVar(var);
- if (!mapped.same_as(var) || defined_.count(var.get())) return mapped;
+ if (auto mapped = GetVarRemap(var)) return mapped.value();
}
return DefineVar(var);
}
- Var mapped = GetRemappedVar(var);
+ Var mapped = GetVarRemap(var).value_or(var);
if (!mapped.same_as(var)) return mapped;
return ffi::Unchanged();
}
@@ -133,10 +136,6 @@ UnchangedOr<PrimExpr> IRConvertSSA::Mutate_(const
prim::LetNode* op, InplaceMode
return result;
}
-Stmt IRConvertSSA::WithScope(const std::function<Stmt()>& body) {
- return scope_.WithNewScope(body);
-}
-
Var IRConvertSSA::DefineVar(Var var) {
Type type = WithDefRegionKind(kTVMFFIDefRegionKindNone, [&] {
return Mutate(var->ty, InplaceMode::kDisallow)
@@ -146,17 +145,15 @@ Var IRConvertSSA::DefineVar(Var var) {
Var result = defined_.count(var.get()) ? MakeNewVar(var) : var;
defined_.insert(var.get());
if (!type.same_as(result->ty)) result = result.CopyWithType(type);
- if (!result.same_as(var)) PushVarRemap(var, result);
+ SetVarRemap(var, result);
return result;
}
-Var IRConvertSSA::GetRemappedVar(Var var) {
- if (auto it = scoped_var_remap_.find(var.get());
- it != scoped_var_remap_.end() && it->second.size()) {
- return it->second.back();
- } else {
- return var;
+ffi::Optional<Var> IRConvertSSA::GetVarRemap(const Var& var) {
+ for (auto it = scope_.rbegin(); it != scope_.rend(); ++it) {
+ if (auto found = it->find(var.get()); found != it->end()) return
found->second;
}
+ return std::nullopt;
}
UnchangedOr<Stmt> IRConvertSSA::Mutate_(const BindNode* op, InplaceMode
inplace_mode) {
@@ -165,7 +162,7 @@ UnchangedOr<Stmt> IRConvertSSA::Mutate_(const BindNode* op,
InplaceMode inplace_
// The ordinary Bind path rewrites the RHS before defining the Var and
propagates its type.
Stmt stmt = StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
Var bound_var = stmt.as<BindNode>()->var;
- if (!bound_var.same_as(GetRemappedVar(var))) PushVarRemap(var, bound_var);
+ if (!bound_var.same_as(GetVarRemap(var).value_or(var))) SetVarRemap(var,
bound_var);
VarRemapSet(var, previous_remap);
return stmt;
}
@@ -196,7 +193,7 @@ UnchangedOr<Stmt> IRConvertSSA::Mutate_(const ForNode* op,
InplaceMode inplace_m
if (defined_.count(v.get())) {
return scope_.WithNewScope([&]() -> Stmt {
Var new_var = MakeNewVar(v);
- PushVarRemap(v, new_var);
+ SetVarRemap(v, new_var);
Stmt stmt =
StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
auto n = ffi::make_object<ForNode>(*stmt.as<ForNode>());
@@ -206,6 +203,7 @@ UnchangedOr<Stmt> IRConvertSSA::Mutate_(const ForNode* op,
InplaceMode inplace_m
} else {
defined_.insert(v.get());
return scope_.WithNewScope([&]() -> Stmt {
+ SetVarRemap(v, v);
return StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
});
}
@@ -239,24 +237,8 @@ Var IRConvertSSA::MakeNewVar(const Var& old_var) {
return Var(old_var->name, old_var->ty, old_var->loc);
}
-void IRConvertSSA::PushVarRemap(const Var& old_var, const Var& new_var) {
- scoped_var_remap_[old_var.get()].push_back(new_var);
- auto& level = scope_.Current();
- level.parent = this;
- level.push_back({old_var, new_var});
-}
-
-void IRConvertSSA::PopAllRemapsInCurrentScope() {
- auto& current = scope_.Current();
- while (current.size()) {
- auto& remap = current.back();
- scoped_var_remap_[remap.old_var.get()].pop_back();
- current.pop_back();
- }
-}
-
-Stmt ConvertSSA(Stmt stmt) {
- return ffi::make_object<IRConvertSSA>()->Mutate(stmt,
InplaceMode::kAllow).ValueOrUnchanged(stmt);
+void IRConvertSSA::SetVarRemap(const Var& old_var, const Var& new_var) {
+ scope_.Current().insert_or_assign(old_var.get(), new_var);
}
IRModule IRConvertSSA::VisitIRModule(IRModule mod) {
@@ -277,6 +259,11 @@ IRModule IRConvertSSA::VisitIRModule(IRModule mod) {
}
return mod;
}
+} // namespace
+
+Stmt ConvertSSA(Stmt stmt) {
+ return ffi::make_object<IRConvertSSA>()->Mutate(stmt,
InplaceMode::kAllow).ValueOrUnchanged(stmt);
+}
namespace transform {
Pass ConvertSSA() {
diff --git a/src/tirx/transform/ir_utils.h b/src/tirx/transform/ir_utils.h
index 614fe5e7e9..63acf59af9 100644
--- a/src/tirx/transform/ir_utils.h
+++ b/src/tirx/transform/ir_utils.h
@@ -27,7 +27,6 @@
#include <tvm/ir/function.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/ir/prim/op.h>
-#include <tvm/ir/scope_stack.h>
#include <tvm/ir/with_context.h>
#include <tvm/runtime/device_api.h>
#include <tvm/sym/int_set.h>
@@ -35,13 +34,11 @@
#include <tvm/tirx/layout.h>
#include <tvm/tirx/op/abi.h>
#include <tvm/tirx/op/memory.h>
-#include <tvm/tirx/stmt_functor.h>
#include <functional>
#include <optional>
#include <string>
#include <unordered_map>
-#include <unordered_set>
#include <utility>
#include <vector>
@@ -92,88 +89,6 @@ inline int GetTempAllocaAlignment(const PrimType& type,
int64_t const_size) {
*/
Stmt ConvertSSA(Stmt stmt);
-/*! \brief Shared SSA renaming algorithm; dialects extend statement dispatch
explicitly. */
-class IRConvertSSA : public StmtExprMutator {
- public:
- TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(IRConvertSSA, StmtExprMutator)
- using StmtExprMutator::Mutate;
- using StmtExprMutator::Mutate_;
- Function VisitFunction(Function func);
- IRModule VisitIRModule(IRModule mod);
-
- protected:
- explicit IRConvertSSA(const VTable* table) : StmtExprMutator(table) {}
- UnchangedOr<Expr> Mutate_(const VarNode* op, InplaceMode inplace_mode) final;
- Stmt WithScope(const std::function<Stmt()>& body);
- Var DefineVar(Var var);
- Var GetRemappedVar(Var var);
- UnchangedOr<Stmt> Mutate_(const BindNode* op, InplaceMode inplace_mode)
final;
- UnchangedOr<Stmt> Mutate_(const IfNode* op, InplaceMode inplace_mode) final;
- UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final;
- UnchangedOr<Stmt> Mutate_(const WhileNode* op, InplaceMode inplace_mode)
final;
- UnchangedOr<Stmt> Mutate_(const RegionStmtNode* op, InplaceMode
inplace_mode) final;
- UnchangedOr<PrimExpr> Mutate_(const prim::LetNode* op, InplaceMode
inplace_mode) final;
- static Var MakeNewVar(const Var& old_var);
- void PushVarRemap(const Var& old_var, const Var& new_var);
- void PopAllRemapsInCurrentScope();
-
- private:
- struct VarRemap {
- Var old_var;
- Var new_var;
- };
- /*! \brief Scope stack: each scope level holds the remaps introduced in that
scope.
- *
- * When a body-carrying statement (For, Allocate, or a dialect statement)
calls
- * scope_.WithNewScope([&]{...}), a new scope level is pushed.
- * Bind statements push their remaps to the current scope.
- * On scope exit, the destructor of std::vector<VarRemap> triggers,
- * and we undo all remaps in that level.
- *
- * Note: ScopeStack<T>::WithNewScope calls T's destructor on exit.
- * std::vector's destructor destroys elements but does NOT call custom
- * cleanup. So we wrap the vector in ScopeLevel which handles cleanup.
- */
- struct ScopeLevel {
- std::vector<VarRemap> remaps;
- IRConvertSSA* parent{nullptr};
-
- void push_back(VarRemap remap) { remaps.push_back(std::move(remap)); }
- size_t size() const { return remaps.size(); }
- VarRemap& back() { return remaps.back(); }
- void pop_back() { remaps.pop_back(); }
-
- void Clear() {
- if (!parent) return;
- while (!remaps.empty()) {
- parent->scoped_var_remap_[remaps.back().old_var.get()].pop_back();
- remaps.pop_back();
- }
- }
- ~ScopeLevel() { Clear(); }
- ScopeLevel() = default;
- ScopeLevel(const ScopeLevel&) = delete;
- ScopeLevel& operator=(const ScopeLevel&) = delete;
- ScopeLevel(ScopeLevel&& other) noexcept
- : remaps(std::move(other.remaps)), parent(other.parent) {
- other.parent = nullptr;
- }
- ScopeLevel& operator=(ScopeLevel&& other) noexcept {
- if (this != &other) {
- Clear();
- remaps = std::move(other.remaps);
- parent = other.parent;
- other.parent = nullptr;
- }
- return *this;
- }
- };
-
- std::unordered_map<const VarNode*, std::vector<Var>> scoped_var_remap_;
- std::unordered_set<const VarNode*> defined_;
- ScopeStack<ScopeLevel> scope_;
-};
-
} // namespace tirx
} // namespace tvm
#endif // TVM_TIR_TRANSFORM_IR_UTILS_H_
