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 9a044debd1 [REFACTOR][TIRx] Use structural definition renewal in core
transforms (#20601)
9a044debd1 is described below
commit 9a044debd110c6bc84c2bd10a6065127c2f44fd0
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Oct 8 10:58:36 2026 -0400
[REFACTOR][TIRx] Use structural definition renewal in core transforms
(#20601)
Provide TIRx-owned definition renewal through structural mapping,
preserving bound and free identities through native structural hooks.
Mark match-buffer targets as pattern definitions and migrate renewal
consumers to the shared helper, removing the S-TIR-specific renewal API.
Remove the S-TIR block precondition from index narrowing.
---
include/tvm/s_tir/stmt.h | 3 +-
include/tvm/s_tir/transform.h | 9 --
include/tvm/tirx/function.h | 11 ++
python/tvm/s_tir/__init__.py | 18 ---
python/tvm/tirx/__init__.py | 2 +-
python/tvm/tirx/function.py | 19 +++
.../transform/lower_global_view_to_local_view.cc | 4 +-
.../transform/attach_attr_layout_free_buffers.cc | 4 +-
src/relax/transform/compute_prim_value.cc | 4 +-
src/relax/transform/fuse_tir.cc | 6 +-
src/relax/transform/split_call_tir_by_pattern.cc | 6 +-
.../transform/split_layout_rewrite_preproc.cc | 6 +-
src/s_tir/stmt.cc | 6 +-
src/s_tir/transform/renew_defs.cc | 175 ---------------------
src/tirx/ir/function.cc | 29 ++++
src/tirx/transform/bind_target.cc | 5 +-
src/tirx/transform/force_narrow_index_to_i32.cc | 8 -
tests/python/s_tir/test_s_tir_renew_defs.py | 14 +-
.../test_tirx_force_narrow_reject_blocks.py | 35 -----
19 files changed, 89 insertions(+), 275 deletions(-)
diff --git a/include/tvm/s_tir/stmt.h b/include/tvm/s_tir/stmt.h
index 16096291d5..817a1008ba 100644
--- a/include/tvm/s_tir/stmt.h
+++ b/include/tvm/s_tir/stmt.h
@@ -57,7 +57,8 @@ class MatchBufferRegionNode : public ffi::Object {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<MatchBufferRegionNode>()
- .def_ro("buffer", &MatchBufferRegionNode::buffer,
refl::AttachFieldFlag::SEqHashDefSimple())
+ .def_ro("buffer", &MatchBufferRegionNode::buffer,
+ refl::AttachFieldFlag::SEqHashDefPattern())
.def_ro("source", &MatchBufferRegionNode::source);
}
diff --git a/include/tvm/s_tir/transform.h b/include/tvm/s_tir/transform.h
index d2aa71098d..5a9c0b7277 100644
--- a/include/tvm/s_tir/transform.h
+++ b/include/tvm/s_tir/transform.h
@@ -35,15 +35,6 @@
namespace tvm {
namespace s_tir {
-/*!
- * \brief Renew the definition nodes for a TIR, including Var, Buffer and
IterVar.
- * This pass works as a simple DeepCopy to duplicate a function with
different Vars and
- * Buffers but the same behavior
- * \param func The input Function.
- * \return The renewed func.
- */
-TVM_DLL tirx::Function RenewDefs(const tirx::Function& func);
-
namespace transform {
using tirx::transform::CreateFunctionPass;
diff --git a/include/tvm/tirx/function.h b/include/tvm/tirx/function.h
index b99547ac27..d84eab6e9e 100644
--- a/include/tvm/tirx/function.h
+++ b/include/tvm/tirx/function.h
@@ -126,6 +126,17 @@ class Function : public BaseFunc {
TVM_DEFINE_OBJECT_REF_COW_METHOD(FunctionNode);
};
+/*!
+ * \brief Renew variable definitions and remap their uses in a function.
+ *
+ * Definition regions follow the registered structural hooks. Free variables
+ * and function attributes are preserved, and the function signature is
rebuilt.
+ *
+ * \param func The function to renew.
+ * \return The function with fresh variable definitions.
+ */
+TVM_DLL Function RenewDef(Function func);
+
/*!
* \brief Specialize parameters of Function.
* \param func The Function to be specialized.
diff --git a/python/tvm/s_tir/__init__.py b/python/tvm/s_tir/__init__.py
index d4c79f0e44..da34ed52b9 100644
--- a/python/tvm/s_tir/__init__.py
+++ b/python/tvm/s_tir/__init__.py
@@ -45,24 +45,6 @@ if not _RUNTIME_ONLY:
from . import dlight
-def renew_defs(func):
- """Re-generate the definition nodes for a TIR, including VarDef, BufferDef.
- This pass works as a simple DeepCopy to duplicate a function with
different Vars and
- Buffers but the same behavior
-
- Parameters
- ----------
- func: Function
- The input function
-
- Returns
- -------
- result : Function
- The new generated func.
- """
- return _ffi_api.RenewDefs(func)
-
-
def _initialize_script_namespace():
from . import script
diff --git a/python/tvm/tirx/__init__.py b/python/tvm/tirx/__init__.py
index 7c100bffee..217fee603d 100644
--- a/python/tvm/tirx/__init__.py
+++ b/python/tvm/tirx/__init__.py
@@ -56,7 +56,7 @@ from .stmt import BufferRegion, BufferRegionType
from .stmt import ScopeIdDefStmt
from .tile_primitive import DispatchContext, TilePrimitiveCall
-from .function import Function, IndexMap
+from .function import Function, IndexMap, renew_def
from .op import call_packed_lowered, call_cpacked_lowered,
register_intrin_lowering
from .op import call_packed, call_cpacked, call_intrin, call_pure_extern,
call_extern
diff --git a/python/tvm/tirx/function.py b/python/tvm/tirx/function.py
index fe081387ac..d5db3c9a3b 100644
--- a/python/tvm/tirx/function.py
+++ b/python/tvm/tirx/function.py
@@ -169,6 +169,25 @@ class Function(BaseFunc, Scriptable):
return _ffi_api.Specialize(self, param_map) # type: ignore
+def renew_def(func: Function) -> Function:
+ """Renew variable definitions and remap their uses in a function.
+
+ Definition regions follow the registered structural hooks. Free variables
+ and function attributes are preserved, and the function signature is
rebuilt.
+
+ Parameters
+ ----------
+ func : Function
+ The function to renew.
+
+ Returns
+ -------
+ Function
+ The function with fresh variable definitions.
+ """
+ return _ffi_api.RenewDef(func)
+
+
@tvm_ffi.register_object("tirx.IndexMap")
class IndexMap(Object):
"""A mapping from multi-dimensional indices to another set of
multi-dimensional indices
diff --git a/src/relax/distributed/transform/lower_global_view_to_local_view.cc
b/src/relax/distributed/transform/lower_global_view_to_local_view.cc
index 0409288048..a6765eba55 100644
--- a/src/relax/distributed/transform/lower_global_view_to_local_view.cc
+++ b/src/relax/distributed/transform/lower_global_view_to_local_view.cc
@@ -29,7 +29,7 @@
#include <tvm/relax/op/ccl.h>
#include <tvm/s_tir/stmt.h>
#include <tvm/s_tir/stmt_functor.h>
-#include <tvm/s_tir/transform.h>
+#include <tvm/tirx/function.h>
#include "../../../s_tir/schedule/transform.h"
#include "utils.h"
@@ -129,7 +129,7 @@ class DistributedBufferCompactor : public
s_tir::StmtExprMutator {
public:
static std::tuple<tirx::Function, std::string> DistBufferCompact(
const std::vector<ShardingSpec>& sharding_specs, tirx::Function
function) {
- function = s_tir::RenewDefs(function);
+ function = tirx::RenewDef(function);
auto compactor =
ffi::make_object<DistributedBufferCompactor>(sharding_specs, function);
ffi::Array<Var> new_params;
ffi::Map<TensorVar, TensorVar> replace_buffer_map;
diff --git a/src/relax/transform/attach_attr_layout_free_buffers.cc
b/src/relax/transform/attach_attr_layout_free_buffers.cc
index 5122b69856..f784cdb617 100644
--- a/src/relax/transform/attach_attr_layout_free_buffers.cc
+++ b/src/relax/transform/attach_attr_layout_free_buffers.cc
@@ -24,7 +24,7 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/relax/expr_functor.h>
#include <tvm/relax/transform.h>
-#include <tvm/s_tir/transform.h>
+#include <tvm/tirx/function.h>
#include <tvm/tirx/stmt_functor.h>
namespace tvm {
@@ -84,7 +84,7 @@ class AttrAttacher : public ExprMutator {
tirx::Function func =
WithAttr(mod_->Lookup(gv).as_or_throw<tirx::Function>(),
"layout_free_buffers", layout_free_buffers);
// Renew defs
- func = s_tir::RenewDefs(func);
+ func = tirx::RenewDef(func);
// Add the updated tirx::Function in the IRModule
// Note the blockbuilder would automatically combine the same tirx function
// So we don't need to worry about the duplicate insertion
diff --git a/src/relax/transform/compute_prim_value.cc
b/src/relax/transform/compute_prim_value.cc
index a7a3590306..2ba822a2b0 100644
--- a/src/relax/transform/compute_prim_value.cc
+++ b/src/relax/transform/compute_prim_value.cc
@@ -23,8 +23,8 @@
#include <tvm/relax/op_attr_types.h>
#include <tvm/relax/transform.h>
#include <tvm/s_tir/analysis.h>
-#include <tvm/s_tir/transform.h>
#include <tvm/tirx/analysis.h>
+#include <tvm/tirx/function.h>
#include <tvm/tirx/op.h>
#include <tvm/tirx/stmt_functor.h>
@@ -116,7 +116,7 @@ class PrimExprComputeInjector : public ExprMutator {
tirx::Function func(param_vars, tirx::SeqStmt(body), ret_ty,
DictAttrs({{tirx::attr::kIsHostFunc, true},
{tvm::attr::kSTir, true}}));
- func = s_tir::RenewDefs(func);
+ func = tirx::RenewDef(func);
auto callee = builder_->AddFunction(func, "compute_symbolic_expr");
diff --git a/src/relax/transform/fuse_tir.cc b/src/relax/transform/fuse_tir.cc
index d73558f0bd..7fdd975c6e 100644
--- a/src/relax/transform/fuse_tir.cc
+++ b/src/relax/transform/fuse_tir.cc
@@ -26,7 +26,7 @@
#include <tvm/relax/type.h>
#include <tvm/s_tir/stmt.h>
#include <tvm/s_tir/stmt_functor.h>
-#include <tvm/s_tir/transform.h>
+#include <tvm/tirx/function.h>
#include <tvm/tirx/op.h>
#include <unordered_map>
@@ -570,7 +570,7 @@ class FusedTIRConstructor : public ExprVisitor {
tirx::Function function_ = mod_->Lookup(gv).as_or_throw<tirx::Function>();
// Step 2. Renew all vars/buffer definitions and blocks to avoid
duplication
- tirx::Function function = s_tir::RenewDefs(function_);
+ tirx::Function function = tirx::RenewDef(function_);
// Step 3. Check functions are all schedulable funcs. i.e. the body of
func is root block
// TODO(Siyuan): support un-schedulable functions.
@@ -889,7 +889,7 @@ class FusedTIRConstructor : public ExprVisitor {
});
tirx::Function func(params, tirx::SeqStmt(body), VoidType(),
DictAttrs(attr_map));
// Renew function defs to prevent using the same symbolic vars in
different functions
- return s_tir::RenewDefs(func);
+ return tirx::RenewDef(func);
}
/*! \brief Get DynTensor numbers from recursive Tuples. */
diff --git a/src/relax/transform/split_call_tir_by_pattern.cc
b/src/relax/transform/split_call_tir_by_pattern.cc
index 50a49adbe7..7be76c2c46 100644
--- a/src/relax/transform/split_call_tir_by_pattern.cc
+++ b/src/relax/transform/split_call_tir_by_pattern.cc
@@ -30,8 +30,8 @@
#include <tvm/relax/type.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.h>
#include "../../s_tir/schedule/ir_comparator.h"
@@ -760,8 +760,8 @@ class SplitMutator : public ExprMutator {
new_call->args = {lib_func, call->args[1]};
return Call(new_call);
}
- tirx::Function func1 = s_tir::RenewDefs(split_funcs.first);
- tirx::Function func2 = s_tir::RenewDefs(split_funcs.second.value());
+ tirx::Function func1 = tirx::RenewDef(split_funcs.first);
+ tirx::Function func2 = tirx::RenewDef(split_funcs.second.value());
TVM_FFI_ICHECK(arg_partition.size() == 2);
// emit the first call to the library kernel
ffi::Array<Expr> args1;
diff --git a/src/relax/transform/split_layout_rewrite_preproc.cc
b/src/relax/transform/split_layout_rewrite_preproc.cc
index 0340b81897..6e6600058d 100644
--- a/src/relax/transform/split_layout_rewrite_preproc.cc
+++ b/src/relax/transform/split_layout_rewrite_preproc.cc
@@ -27,7 +27,7 @@
#include <tvm/relax/transform.h>
#include <tvm/s_tir/stmt.h>
#include <tvm/s_tir/stmt_functor.h>
-#include <tvm/s_tir/transform.h>
+#include <tvm/tirx/function.h>
#include <algorithm>
#include <cstddef>
@@ -101,7 +101,7 @@ class SplitFunctionLayoutRewrite : public
s_tir::StmtExprMutator {
DictAttrs attrs(dict);
tirx::Function func = tirx::Function(params, SeqStmt(body), VoidType(),
attrs);
- return s_tir::RenewDefs(func);
+ return tirx::RenewDef(func);
}
tirx::Function create_compute_func() const {
@@ -147,7 +147,7 @@ class SplitFunctionLayoutRewrite : public
s_tir::StmtExprMutator {
DictAttrs attrs(dict);
tirx::Function func = tirx::Function(params, SeqStmt(body), VoidType(),
attrs);
- return s_tir::RenewDefs(func);
+ return tirx::RenewDef(func);
}
void visit_root_block(const s_tir::SBlockNode* op) {
diff --git a/src/s_tir/stmt.cc b/src/s_tir/stmt.cc
index 711b64ba66..d86b9ae42d 100644
--- a/src/s_tir/stmt.cc
+++ b/src/s_tir/stmt.cc
@@ -83,7 +83,7 @@ TVM_FFI_INLINE
ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> MatchBufferRegi
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
MatchBufferRegionNode>(
value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->WithDefRegionKind(
- kTVMFFIDefRegionKindSimple, [&]() { return
visitor->VisitExpected(self->buffer); }));
+ kTVMFFIDefRegionKindPattern, [&]() { return
visitor->VisitExpected(self->buffer); }));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->source));
return std::nullopt;
}
@@ -94,7 +94,7 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>>
MatchBufferRegionMutate
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
MatchBufferRegionNode>(
value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<TensorVar>, mapped_buffer,
-
mutator->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]() {
+
mutator->WithDefRegionKind(kTVMFFIDefRegionKindPattern, [&]() {
return
mutator->MutateExpected(self->buffer);
}));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<TensorRegion>,
mapped_source,
@@ -115,7 +115,7 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>>
MatchBufferRegionMaybeI
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
MatchBufferRegionNode>(
value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<TensorVar>, mapped_buffer,
-
mutator->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]() {
+
mutator->WithDefRegionKind(kTVMFFIDefRegionKindPattern, [&]() {
return
mutator->MutateExpected(self->buffer,
ffi::InplaceMode::kAllow);
}));
diff --git a/src/s_tir/transform/renew_defs.cc
b/src/s_tir/transform/renew_defs.cc
deleted file mode 100644
index 71636ed46f..0000000000
--- a/src/s_tir/transform/renew_defs.cc
+++ /dev/null
@@ -1,175 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-/*!
- * \file renew_defs.cc
- * \brief Renew the definition nodes for a TIR, including Var, Buffer and
IterVar.
- */
-
-#include <tvm/ffi/cast.h>
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/s_tir/stmt.h>
-#include <tvm/s_tir/stmt_functor.h>
-#include <tvm/s_tir/transform.h>
-
-namespace tvm {
-namespace s_tir {
-using namespace tvm::tirx;
-
-#define STMT_REGENERATE_VAR_DEF(NODE, FIELD)
\
- UnchangedOr<Stmt> Mutate_(const NODE* op, InplaceMode inplace_mode) final {
\
- Var new_var = this->ReDefineVar(op->FIELD);
\
- Stmt stmt =
\
- StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op)); \
- op = stmt.as<NODE>();
\
- TVM_FFI_ICHECK(op != nullptr);
\
- auto n = ffi::make_object<NODE>(*op);
\
- n->FIELD =
std::move(new_var).as_or_throw<std::decay_t<decltype(n->FIELD)>>(); \
- return Stmt(n);
\
- }
-
-class RenewDefMutator : public StmtExprMutator {
- public:
- using StmtExprMutator::Mutate;
- using StmtExprMutator::Mutate_;
-
- static Function Transform(const Function& func) {
- auto generator = ffi::make_object<RenewDefMutator>();
- // Establish explicit parameters and symbols in their types before
visiting uses.
- for (const Var& param : func->params) {
- if (param->ty.as<PrimTypeNode>()) generator->ReDefineVar(param);
- }
- ffi::Array<Var> params = func->params.Map([&](const Var& param) {
- auto mapped = generator->VarRemapGet(param);
- return mapped != nullptr ? mapped.as_or_throw<Var>() :
generator->ReDefineVar(param);
- });
- // Visit body
- auto body = generator->Mutate(func->body).ValueOrUnchanged(func->body);
- // Recreate function
- return Function(params, body, func->ret_type, func->attrs, func->span);
- }
-
- private:
- UnchangedOr<Stmt> Mutate_(const BindNode* op, InplaceMode inplace_mode)
final {
- // Result-type patterns may introduce metadata symbols. Renew those before
- // rewriting the RHS; the Bind's own definition is published only
afterwards.
- WithDefRegionKind(kTVMFFIDefRegionKindPattern,
- [&] { return Mutate(op->value->ty,
InplaceMode::kDisallow); });
- return StmtExprMutator::Mutate_(op, inplace_mode);
- }
-
- STMT_REGENERATE_VAR_DEF(ForNode, loop_var);
-
- UnchangedOr<Expr> Mutate_(const VarNode* op, InplaceMode inplace_mode) final
{
- Var var = ffi::GetRef<Var>(op);
- if (auto mapped = VarRemapGet(var); mapped != nullptr) {
- return mapped.as_or_throw<Expr>();
- }
- if (def_region_kind() != kTVMFFIDefRegionKindNone) return ReDefineVar(var);
- return ffi::Unchanged();
- }
-
- UnchangedOr<Stmt> Mutate_(const SBlockNode* op, InplaceMode inplace_mode)
final {
- // Step 0. Re-define Itervars
- ffi::Array<IterVar> iter_vars =
- op->iter_vars.Map(std::bind(&RenewDefMutator::VisitIterVar, this,
std::placeholders::_1));
-
- // Step 1. Re-define buffers allocated under the block
- ffi::Array<TensorVar> alloc_buffers = op->alloc_buffers.Map([this](const
TensorVar& buf) {
- return this->ReDefineVar(buf.var()).as_or_throw<TensorVar>();
- });
-
- // Step 2. Re-define match_buffers
- ffi::Array<MatchBufferRegion> match_buffers = op->match_buffers.Map(
- std::bind(&RenewDefMutator::VisitMatchBuffer, this,
std::placeholders::_1));
-
- // Step 3. Visit body
- ffi::Optional<SeqStmt> init = std::nullopt;
- if (op->init.has_value()) {
- init = this->Mutate(op->init.value(),
inplace_mode).ValueOrUnchanged(op->init.value());
- }
- Stmt body = this->Mutate(op->body,
inplace_mode).ValueOrUnchanged(op->body);
-
- // Step 4. Revisit access region
- ffi::Array<TensorRegion> reads =
- op->reads.Map(std::bind(&RenewDefMutator::VisitBufferRegion, this,
std::placeholders::_1));
- ffi::Array<TensorRegion> writes =
- op->writes.Map(std::bind(&RenewDefMutator::VisitBufferRegion, this,
std::placeholders::_1));
-
- // Step 5. Regenerate block. Since the defs are changed, we need to create
a new block
- auto n = ffi::make_object<SBlockNode>(*op);
- n->iter_vars = std::move(iter_vars);
- n->alloc_buffers = std::move(alloc_buffers);
- n->match_buffers = std::move(match_buffers);
- n->reads = std::move(reads);
- n->writes = std::move(writes);
- n->body = std::move(body);
- n->init = std::move(init);
-
- return Stmt(n);
- }
-
- Var ReDefineVar(const Var& var) {
- if (auto mapped = VarRemapGet(var); mapped != nullptr) return
mapped.as_or_throw<Var>();
- Type type = WithDefRegionKind(kTVMFFIDefRegionKindPattern, [&] {
- return Mutate(var->ty, InplaceMode::kDisallow)
- .as_or_throw<UnchangedOr<Type>>()
- .ValueOrUnchanged(var->ty);
- });
- Var new_var(var->name, type, var->span);
- VarRemapSet(var, new_var);
- return new_var;
- }
-
- IterVar VisitIterVar(const IterVar& iter_var) {
- auto mapped = VarRemapGet(iter_var);
- if (mapped != nullptr) return mapped.as_or_throw<IterVar>();
- PrimExpr min =
- Mutate(iter_var->dom->min,
InplaceMode::kDisallow).ValueOrUnchanged(iter_var->dom->min);
- PrimExpr extent = Mutate(iter_var->dom->extent, InplaceMode::kDisallow)
- .ValueOrUnchanged(iter_var->dom->extent);
- IterVar new_iter_var(Range::FromMinExtent(min, extent),
- ReDefineVar(iter_var->var).as_or_throw<PrimVar>(),
iter_var->iter_type,
- iter_var->thread_tag);
- VarRemapSet(iter_var, new_iter_var);
- return new_iter_var;
- }
-
- MatchBufferRegion VisitMatchBuffer(const MatchBufferRegion& match_buffer) {
- TensorVar buffer =
ReDefineVar(match_buffer->buffer.var()).as_or_throw<TensorVar>();
- TensorRegion region = VisitBufferRegion(match_buffer->source);
- return MatchBufferRegion(std::move(buffer), std::move(region));
- }
-
- TensorRegion VisitBufferRegion(const TensorRegion& buffer_region) {
- return Mutate(buffer_region, InplaceMode::kDisallow)
- .ValueOrUnchanged(buffer_region)
- .as_or_throw<TensorRegion>();
- }
-};
-
-Function RenewDefs(const Function& func) { return
RenewDefMutator::Transform(func); }
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("s_tir.RenewDefs", RenewDefs);
-}
-
-} // namespace s_tir
-} // namespace tvm
diff --git a/src/tirx/ir/function.cc b/src/tirx/ir/function.cc
index 97db05bc86..f6ad45e077 100644
--- a/src/tirx/ir/function.cc
+++ b/src/tirx/ir/function.cc
@@ -123,6 +123,34 @@ Function::Function(ffi::Array<tirx::Var> params,
ffi::Optional<SeqStmt> body, Ty
data_ = std::move(n);
}
+Function RenewDef(Function func) {
+ // Seed original definitions so identities created by structural hooks are
preserved.
+ ffi::Map<Var, Var> definition_remap;
+ ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(func,
+ [&](const Var& var,
TVMFFIDefRegionKind kind) {
+ if (kind !=
kTVMFFIDefRegionKindNone)
+
definition_remap.Set(var, var);
+ return
ffi::WalkResult::Advance();
+ });
+ return ffi::StructuralMap<ffi::WalkOrder::kPostOrder>(
+ func,
+ [remap = std::move(definition_remap)](const Var& var,
+ TVMFFIDefRegionKind kind)
mutable {
+ auto mapped = remap.Get(var);
+ if (!mapped.has_value()) return var;
+ if (!mapped.value().same_as(var)) return mapped.value();
+ if (kind == kTVMFFIDefRegionKindNone) return var;
+ Var fresh(var->name, var->ty, var->span);
+ remap.Set(var, fresh);
+ return fresh;
+ },
+ [](const Function& mapped) {
+ return Function(mapped->params, mapped->body, mapped->ret_type,
mapped->attrs,
+ mapped->span);
+ })
+ .as_or_throw<Function>();
+}
+
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
FunctionNode::RegisterReflection();
@@ -137,6 +165,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
"tirx.Function",
[](ffi::Array<tirx::Var> params, ffi::Optional<SeqStmt> body, Type
ret_type, DictAttrs attrs,
Span span) { return Function(params, body, ret_type, attrs, span); });
+ refl::GlobalDef().def("tirx.RenewDef", RenewDef);
}
FuncType FunctionNode::func_type_annotation() const {
diff --git a/src/tirx/transform/bind_target.cc
b/src/tirx/transform/bind_target.cc
index 5379904f0e..b36a32e8f1 100644
--- a/src/tirx/transform/bind_target.cc
+++ b/src/tirx/transform/bind_target.cc
@@ -37,8 +37,7 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/unique_name_supply.h>
-#include <tvm/s_tir/stmt.h>
-#include <tvm/s_tir/transform.h>
+#include <tvm/tirx/function.h>
#include <tvm/tirx/op.h>
#include <tvm/tirx/stmt_functor.h>
#include <tvm/tirx/transform.h>
@@ -308,7 +307,7 @@ IRModule BindTarget(IRModule mod, const Target& target) {
if (called_by_host && called_by_device) {
// Rule 4.1: Called by both host and device
// Bind device target to current function
- Function host_func = s_tir::RenewDefs(function);
+ Function host_func = tirx::RenewDef(function);
new_mod->Update(gvar,
WithAttr(std::move(function), tvm::attr::kTarget,
target_without_host));
diff --git a/src/tirx/transform/force_narrow_index_to_i32.cc
b/src/tirx/transform/force_narrow_index_to_i32.cc
index 1fe5192d00..e48e06242b 100644
--- a/src/tirx/transform/force_narrow_index_to_i32.cc
+++ b/src/tirx/transform/force_narrow_index_to_i32.cc
@@ -27,7 +27,6 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>
-#include <tvm/s_tir/stmt.h>
#include <tvm/tirx/stmt_functor.h>
#include <tvm/tirx/transform.h>
@@ -37,13 +36,6 @@ namespace tirx {
class Int32DTypeNarrower : public
Int32DTypeNarrowerBase<IndexDataTypeNormalizer> {
public:
static Function RewriteDataType(Function func) {
- // The TIRX normalizer does not rewrite S-TIR block iterators, regions, or
match buffers, so
- // narrowing a function that still contains blocks would leave their index
types inconsistent.
- if (func->body.has_value() &&
ContainsNode<s_tir::SBlockRealizeNode>(func->body.value())) {
- TVM_FFI_THROW(ValueError)
- << "tirx.transform.ForceNarrowIndexToInt32 requires a function
without S-TIR blocks. "
- << "Use s_tir.transform.ForceNarrowIndexToInt32 before block
lowering.";
- }
CheckBufferParams(func);
auto narrower = ffi::make_object<Int32DTypeNarrower>(func);
return narrower->Rewrite(func);
diff --git a/tests/python/s_tir/test_s_tir_renew_defs.py
b/tests/python/s_tir/test_s_tir_renew_defs.py
index e4fbfe7a0a..2517a2b50f 100644
--- a/tests/python/s_tir/test_s_tir_renew_defs.py
+++ b/tests/python/s_tir/test_s_tir_renew_defs.py
@@ -66,7 +66,7 @@ def test_simple():
B[vi, vj] = A[vi, vj] * 2.0
f1 = elementwise
- f2 = tvm.s_tir.renew_defs(f1)
+ f2 = tvm.tirx.renew_def(f1)
tvm.ir.assert_structural_equal(f1, f2)
_check_func_signature_remap(f1, f2)
@@ -109,7 +109,7 @@ def test_match_buffer():
B[vi, vj] = A0[vi, vj] * 2.0
f1 = func_match_buffer
- f2 = tvm.s_tir.renew_defs(f1)
+ f2 = tvm.tirx.renew_def(f1)
tvm.ir.assert_structural_equal(f1, f2)
_check_func_signature_remap(f1, f2)
@@ -145,7 +145,7 @@ def test_undefined_buffer():
A[i] = A[i] + T.float16(1.0)
f1 = access_alloc
- f2 = tvm.s_tir.renew_defs(f1)
+ f2 = tvm.tirx.renew_def(f1)
tvm.ir.assert_structural_equal(f1, f2)
# AllocTensor is now a flat statement in SeqStmt
@@ -168,7 +168,7 @@ def test_symbolic_func():
B[i, j * 2 + 1] = A[i, j]
f1 = symbolic_func
- f2 = tvm.s_tir.renew_defs(f1)
+ f2 = tvm.tirx.renew_def(f1)
tvm.ir.assert_structural_equal(f1, f2)
@@ -183,7 +183,7 @@ def test_buffer_params():
B[vi, vj] = A[vi * 2 + vj]
f1 = main
- f2 = tvm.s_tir.renew_defs(main)
+ f2 = tvm.tirx.renew_def(main)
tvm.ir.assert_structural_equal(f1, f2)
assert f1.params[1].shape[0] != f2.params[1].shape[0]
@@ -192,7 +192,7 @@ def test_compound_buffer_param_shape_var():
n = tvm.tirx.Var("n", "int32")
A = tvm.tirx.decl_tensor((tvm.tirx.max(n, 1),), layout=None)
f1 = tvm.tirx.Function([A], tvm.tirx.Evaluate(n))
- f2 = tvm.s_tir.renew_defs(f1)
+ f2 = tvm.tirx.renew_def(f1)
tvm.ir.assert_structural_equal(f1, f2)
assert not f1.body[0].value.same_as(f2.body[0].value)
@@ -214,7 +214,7 @@ def test_gather():
T_take[v_ax0, v_ax1] = A[B[v_ax0], v_ax1]
f1 = take
- f2 = tvm.s_tir.renew_defs(take)
+ f2 = tvm.tirx.renew_def(take)
tvm.ir.assert_structural_equal(f1, f2)
diff --git
a/tests/python/s_tir/transform/test_tirx_force_narrow_reject_blocks.py
b/tests/python/s_tir/transform/test_tirx_force_narrow_reject_blocks.py
deleted file mode 100644
index 6a8db049c6..0000000000
--- a/tests/python/s_tir/transform/test_tirx_force_narrow_reject_blocks.py
+++ /dev/null
@@ -1,35 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-import pytest
-
-import tvm
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
-
-
-def test_reject_blocks():
- @Ts.function(private=True)
- def before(A: T.Tensor((128,), "float32"), B: T.Tensor((128,), "float32")):
- for i in T.serial(0, T.int64(16)):
- for j in T.serial(0, T.int64(8)):
- with Ts.sblock():
- vi = Ts.axis.spatial(T.int64(128), i * T.int64(8) + j)
- B[vi] = A[vi] + T.float32(1)
-
- mod = tvm.IRModule.from_expr(before)
- with pytest.raises(ValueError, match="requires a function without S-TIR
blocks"):
- tvm.tirx.transform.ForceNarrowIndexToInt32()(mod)
