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 ebeebf1116 [REFACTOR][TIRx] Reduce conditional hoisting to HoistIf
(#20636)
ebeebf1116 is described below
commit ebeebf1116fb2360fa74a82e5f07ffaa7916e87b
Author: Tianqi Chen <[email protected]>
AuthorDate: Sat Oct 10 10:30:13 2026 -0400
[REFACTOR][TIRx] Reduce conditional hoisting to HoistIf (#20636)
---
include/tvm/s_tir/function.h | 2 -
include/tvm/s_tir/transform.h | 22 -
include/tvm/tirx/transform.h | 6 +
python/tvm/backend/trn/pipeline.py | 2 +-
python/tvm/s_tir/backend/adreno/pipeline.py | 2 +-
python/tvm/s_tir/pipeline.py | 2 +-
python/tvm/s_tir/transform/transform.py | 33 -
python/tvm/tirx/transform/transform.py | 11 +
src/ir/side_effect.cc | 16 +-
src/s_tir/transform/hoist_expression.cc | 686 ---------------------
src/tirx/transform/hoist_if.cc | 264 ++++++++
src/tirx/transform/remove_no_op.cc | 1 -
.../python/s_tir/dlight/test_gpu_low_batch_gemv.py | 8 +-
.../test_s_tir_transform_hoist_expression.py | 590 ------------------
.../transform/test_s_tir_transform_hoist_if.py | 562 -----------------
.../transform/test_s_tir_transform_thread_sync.py | 25 +
.../tirx/transform/test_transform_hoist_if.py | 423 +++++++++++++
17 files changed, 749 insertions(+), 1906 deletions(-)
diff --git a/include/tvm/s_tir/function.h b/include/tvm/s_tir/function.h
index f75cf7405a..b5044fbf91 100644
--- a/include/tvm/s_tir/function.h
+++ b/include/tvm/s_tir/function.h
@@ -40,8 +40,6 @@ constexpr const char* kLayoutFreeBuffers =
"layout_free_buffers";
constexpr const char* kLayoutTransforms = "layout_transforms";
constexpr const char* kEstimatedFlops = "estimated_flops";
-constexpr const char* kHoistIfThenElseExprWithBlock =
"tirx.HoistIfThenElseExprWithBlock";
-
} // namespace attr
} // namespace s_tir
} // namespace tvm
diff --git a/include/tvm/s_tir/transform.h b/include/tvm/s_tir/transform.h
index cc22793fc3..6f166ab1ef 100644
--- a/include/tvm/s_tir/transform.h
+++ b/include/tvm/s_tir/transform.h
@@ -229,28 +229,6 @@ TVM_DLL Pass LoopPartition();
*/
TVM_DLL Pass InjectVirtualThread();
-/*!
- * \brief Hoist loop-invariant If nodes to
- * outside the eligible loops.
- *
- * \param variant The variant of the pass.
- * variant can have any one of following values ["basic", ""(Default)].
- * \return The pass.
- */
-TVM_DLL Pass HoistIfThenElse(tvm::ffi::String variant = "");
-
-/*!
- * \brief Hoist loop-invariant expressions to outside the eligible loops.
- *
- * Can hoist conditionals used in If statements and
- * expressions, bindings of variables in Let statements and
- * expressions, or boolean expressions, configurable to enable/disable
- * each hoistable type.
- *
- * \return The pass.
- */
-TVM_DLL Pass HoistExpression();
-
/*!
* \brief Detect and rewrite unsafe select that contains memory access.
* \return The pass.
diff --git a/include/tvm/tirx/transform.h b/include/tvm/tirx/transform.h
index 74ba52ade0..6cb8a2d5ec 100644
--- a/include/tvm/tirx/transform.h
+++ b/include/tvm/tirx/transform.h
@@ -98,6 +98,12 @@ TVM_DLL Pass UnrollLoop();
*/
TVM_DLL Pass RemoveNoOp();
+/*!
+ * \brief Hoist loop-invariant If statements outside eligible loops.
+ * \return The pass.
+ */
+TVM_DLL Pass HoistIf();
+
/*!
* \brief Renormalize the split pattern from floordiv(floormod()) to
floormod(floordiv()).
* \return The pass.
diff --git a/python/tvm/backend/trn/pipeline.py
b/python/tvm/backend/trn/pipeline.py
index dd9b252e1c..9d8b6b334b 100644
--- a/python/tvm/backend/trn/pipeline.py
+++ b/python/tvm/backend/trn/pipeline.py
@@ -38,7 +38,7 @@ def trn_pipeline():
tirx.transform.StmtSimplify(),
tvm.s_tir.transform.LoopPartition(),
tirx.transform.LowerTIRxOpaque(),
- tvm.s_tir.transform.HoistIfThenElse(),
+ tirx.transform.HoistIf(),
tirx.transform.StmtSimplify(),
tirx.transform.RemoveNoOp(),
tirx.transform.AnnotateEntryFunc(),
diff --git a/python/tvm/s_tir/backend/adreno/pipeline.py
b/python/tvm/s_tir/backend/adreno/pipeline.py
index 749a245ca5..3d6edad5e3 100644
--- a/python/tvm/s_tir/backend/adreno/pipeline.py
+++ b/python/tvm/s_tir/backend/adreno/pipeline.py
@@ -62,7 +62,7 @@ def default_tir_pipeline(*, prepare_only=False):
passes.append(tirx.transform.StorageRewrite())
passes.extend(
[
- s_tir.transform.HoistIfThenElse(),
+ tirx.transform.HoistIf(),
tirx.transform.UnrollLoop(),
tirx.transform.RenormalizeSplitPattern(),
s_tir.transform.StmtSimplify(),
diff --git a/python/tvm/s_tir/pipeline.py b/python/tvm/s_tir/pipeline.py
index 09e36aeaa7..ea9ee4d2e5 100644
--- a/python/tvm/s_tir/pipeline.py
+++ b/python/tvm/s_tir/pipeline.py
@@ -62,7 +62,7 @@ def default_s_tir_pipeline(*, prepare_only=False):
passes.append(tirx.transform.StorageRewrite())
passes.extend(
[
- s_tir.transform.HoistIfThenElse(),
+ tirx.transform.HoistIf(),
tirx.transform.UnrollLoop(),
tirx.transform.RenormalizeSplitPattern(),
s_tir.transform.StmtSimplify(),
diff --git a/python/tvm/s_tir/transform/transform.py
b/python/tvm/s_tir/transform/transform.py
index 1b4fc61f48..46eedbc80f 100644
--- a/python/tvm/s_tir/transform/transform.py
+++ b/python/tvm/s_tir/transform/transform.py
@@ -257,39 +257,6 @@ def InjectVirtualThread():
return _ffi_api.InjectVirtualThread() # type: ignore
-def HoistIfThenElse(variant=None):
- """Hoist loop-invariant If nodes to outside the eligible loops.
-
- Parameters
- ----------
- variant : Optional[String]
- The variant of the pass.
- variant can have any one of following values ["basic", None(Default)].
-
- Returns
- -------
- fpass : tvm.transform.Pass
- The result pass
- """
- if variant == "basic":
- return _ffi_api.HoistIfThenElseBasic() # type: ignore
- elif variant is None:
- return _ffi_api.HoistIfThenElse() # type: ignore
- else:
- raise ValueError("wrong variant of HoistIfThenElse, " + variant)
-
-
-def HoistExpression():
- """Hoist loop-invariant expressions to outside the eligible loops.
-
- Returns
- -------
- fpass : tvm.transform.Pass
- The result pass
- """
- return _ffi_api.HoistExpression() # type: ignore
-
-
def RewriteUnsafeSelect():
"""Detect and rewrite unsafe select that contains memory access.
diff --git a/python/tvm/tirx/transform/transform.py
b/python/tvm/tirx/transform/transform.py
index 12e69f0551..2909d9107b 100644
--- a/python/tvm/tirx/transform/transform.py
+++ b/python/tvm/tirx/transform/transform.py
@@ -133,6 +133,17 @@ def RemoveNoOp():
return _ffi_api.RemoveNoOp() # type: ignore
+def HoistIf():
+ """Hoist loop-invariant If statements outside eligible loops.
+
+ Returns
+ -------
+ fpass : tvm.transform.Pass
+ The result pass.
+ """
+ return _ffi_api.HoistIf()
+
+
def RemoveAssume():
"""Remove all instances of prim::assume_op
diff --git a/src/ir/side_effect.cc b/src/ir/side_effect.cc
index 7173f29543..6a3bd54418 100644
--- a/src/ir/side_effect.cc
+++ b/src/ir/side_effect.cc
@@ -25,12 +25,21 @@
#include <tvm/ir/op.h>
#include <algorithm>
+#include <unordered_set>
namespace tvm {
CallEffectKind SideEffect(const Expr& expr) {
static auto effects =
Op::GetAttrMap<TCallEffectKind>(tvm::op_attr::kCallEffectKind);
CallEffectKind kind = CallEffectKind::kPure;
+ std::unordered_set<const ExprNode*> visited;
+ auto visit_expr =
+ [&](const ExprNode* node,
+ ffi::StructuralVisitorObj* visitor) ->
ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> {
+ // Effects are context-independent; inspect each shared expression only
once.
+ if (!visited.insert(node).second) return std::nullopt;
+ return visitor->DefaultVisitExpected(node);
+ };
ffi::StructuralVisit(
expr,
[&](const CallNode* node,
@@ -39,18 +48,19 @@ CallEffectKind SideEffect(const Expr& expr) {
effects.get(node->op,
static_cast<TCallEffectKind>(CallEffectKind::kOpaque)));
kind = std::max(kind, std::min(effect, CallEffectKind::kUpdateState));
if (kind == CallEffectKind::kUpdateState) return ffi::VisitInterrupt();
- return visitor->DefaultVisitExpected(node);
+ return visit_expr(node, visitor);
},
[&](const TensorLoadNode* node,
ffi::StructuralVisitorObj* visitor) ->
ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> {
kind = std::max(kind, CallEffectKind::kReadState);
- return visitor->DefaultVisitExpected(node);
+ return visit_expr(node, visitor);
},
[](const TypeNode*,
ffi::StructuralVisitorObj*) ->
ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> {
// Effects describe evaluated expressions, not expressions embedded in
types.
return std::nullopt;
- });
+ },
+ visit_expr);
return kind;
}
diff --git a/src/s_tir/transform/hoist_expression.cc
b/src/s_tir/transform/hoist_expression.cc
deleted file mode 100644
index d93f2cbe59..0000000000
--- a/src/s_tir/transform/hoist_expression.cc
+++ /dev/null
@@ -1,686 +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 hoist_expression.cc
- */
-#include <tvm/ffi/cast.h>
-#include <tvm/ffi/extra/structural_visit.h>
-#include <tvm/ffi/function.h>
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/analysis.h>
-#include <tvm/ir/prim/expr.h>
-#include <tvm/ir/prim/op.h>
-#include <tvm/s_tir/analysis.h>
-#include <tvm/s_tir/function.h>
-#include <tvm/s_tir/stmt_functor.h>
-#include <tvm/s_tir/transform.h>
-#include <tvm/sym/analyzer.h>
-#include <tvm/tirx/analysis.h>
-#include <tvm/tirx/op/memory.h>
-#include <tvm/tirx/op/region.h>
-
-#include <queue>
-#include <unordered_map>
-#include <unordered_set>
-#include <utility>
-
-#include "../../runtime/thread_storage_scope.h"
-#include "../../s_tir/ir/ir_mutator_with_analyzer.h"
-#include "../../sym/interval_set.h"
-#include "ir_utils.h"
-
-namespace tvm {
-namespace s_tir {
-using namespace tvm::tirx;
-
-enum class HoistedConditionals : int {
- kNone = 0,
- kIfElseStmt = (1 << 0),
- kIfElseExpr = (1 << 1),
- kBooleanExpression = (1 << 2),
- kUsingBlockVar = (1 << 3),
-};
-
-enum class HoistedLetBindings : int {
- kNone = 0,
- kRequiredByCondition = (1 << 0),
- kBind = (1 << 1),
- kLetExpr = (1 << 2),
-};
-
-struct HoistExpressionConfigNode : public ffi::Object {
- int hoisted_conditionals;
- int hoisted_let_bindings;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<HoistExpressionConfigNode>()
- .def_ro("hoisted_conditionals",
&HoistExpressionConfigNode::hoisted_conditionals,
- "Bitflags for the types of boolean expressions to hoist",
-
refl::DefaultValue(static_cast<int>(HoistedConditionals::kIfElseStmt) |
-
static_cast<int>(HoistedConditionals::kIfElseExpr) |
-
static_cast<int>(HoistedConditionals::kBooleanExpression)))
- .def_ro("hoisted_let_bindings",
&HoistExpressionConfigNode::hoisted_let_bindings,
- "Bitflags for the types of let bindings to hoist",
-
refl::DefaultValue(static_cast<int>(HoistedLetBindings::kRequiredByCondition) |
- static_cast<int>(HoistedLetBindings::kBind)
|
-
static_cast<int>(HoistedLetBindings::kLetExpr)));
- }
-
- bool FlagSet(HoistedConditionals flag) const {
- return static_cast<int>(flag) & hoisted_conditionals;
- }
- bool FlagSet(HoistedLetBindings flag) const {
- return static_cast<int>(flag) & hoisted_let_bindings;
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.transform.HoistExpressionConfig",
- HoistExpressionConfigNode, ffi::Object);
-};
-
-class HoistExpressionConfig : public ffi::ObjectRef {
- public:
- HoistExpressionConfig(int hoisted_conditionals, int hoisted_let_bindings) {
- auto node = ffi::make_object<HoistExpressionConfigNode>();
- node->hoisted_conditionals = hoisted_conditionals;
- node->hoisted_let_bindings = hoisted_let_bindings;
- data_ = std::move(node);
- }
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(HoistExpressionConfig,
ffi::ObjectRef,
- HoistExpressionConfigNode);
-};
-
-TVM_FFI_STATIC_INIT_BLOCK() { HoistExpressionConfigNode::RegisterReflection();
}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- ::tvm::transform::PassContext::RegisterConfigOption<HoistExpressionConfig>(
- "s_tir.HoistExpression");
-}
-
-struct HoistIfThenElseConfigNode : public ffi::Object {
- bool support_block_scope_hoisting;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<HoistIfThenElseConfigNode>().def_ro(
- "support_block_scope_hoisting",
&HoistIfThenElseConfigNode::support_block_scope_hoisting,
- "Hoist if cond with block scope variables", refl::DefaultValue(false));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.transform.HoistIfThenElseConfig",
- HoistIfThenElseConfigNode, ffi::Object);
-};
-
-class HoistIfThenElseConfig : public ffi::ObjectRef {
- public:
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(HoistIfThenElseConfig,
ffi::ObjectRef,
- HoistIfThenElseConfigNode);
-};
-
-TVM_FFI_STATIC_INIT_BLOCK() { HoistIfThenElseConfigNode::RegisterReflection();
}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- ::tvm::transform::PassContext::RegisterConfigOption<HoistIfThenElseConfig>(
- "s_tir.HoistIfThenElse");
-}
-
-class HoistInfoCollector : public StmtExprVisitor {
- public:
- using StmtExprVisitor::Visit_;
- struct ConditionInfo {
- ConditionInfo(PrimExpr condition, HoistedConditionals hoist_from, bool
uses_block_var,
- std::unordered_set<const VarNode*> required_let_bindings,
bool generate_else_case)
- : condition(condition),
- hoist_from(hoist_from),
- uses_block_var(uses_block_var),
- required_let_bindings(required_let_bindings),
- generate_else_case(generate_else_case) {}
- PrimExpr condition;
- HoistedConditionals hoist_from;
- bool uses_block_var;
- std::unordered_set<const VarNode*> required_let_bindings;
- bool generate_else_case;
-
- bool IsEnabled(const HoistExpressionConfig& config) const {
- bool valid_source = config->FlagSet(hoist_from);
-
- bool all_required_bindings_are_hoisted =
- required_let_bindings.empty() ||
- config->FlagSet(HoistedLetBindings::kRequiredByCondition) ||
- config->FlagSet(HoistedLetBindings::kBind);
-
- bool valid_block_var_usage =
- config->FlagSet(HoistedConditionals::kUsingBlockVar) ||
!uses_block_var;
- return valid_source && all_required_bindings_are_hoisted &&
valid_block_var_usage;
- }
- };
-
- struct LetBindingInfo {
- LetBindingInfo(Var var, PrimExpr value, HoistedLetBindings hoist_from)
- : var(var), value(value), hoist_from(hoist_from) {}
- Var var;
- PrimExpr value;
- HoistedLetBindings hoist_from;
-
- bool IsEnabled(const HoistExpressionConfig& config) const {
- return config->FlagSet(hoist_from);
- }
- };
-
- struct HoistInfo {
- // The loop variable
- Var loop_var;
-
- // The loop or scope that defines the loop var.
- Stmt loop_def;
-
- // Bindings defined in Bind nodes inside the for-loop whose value
- // does not depend on the loop variable. These can be hoisted
- // outside this for-loop.
- std::vector<LetBindingInfo> let_bindings;
-
- // Conditions evaluated inside the for-loop whose value does not
- // depend on the loop variable. These can be hoisted outside this
- // for loop. These may depend on the let_bindings.
- std::vector<ConditionInfo> conditions;
-
- // Only conditions that impact the entire body of the loop
- // hoisted. Conditionals may not be hoisted from inside a
- // sequential node to outside.
- bool reached_sequential_node{false};
-
- // True if the loop variable representing a block variable
- // (e.g. blockIdx.x, threadIdx.x), false otherwise.
- bool IsBlockVariable() const { return !loop_def.as<ForNode>(); }
- };
-
- static std::vector<HoistInfo> Collect(Stmt stmt, HoistExpressionConfig
config) {
- auto collector = ffi::make_object<HoistInfoCollector>(config);
- collector->Visit(stmt);
- return collector->completed_loops;
- }
-
- private:
- using Parent = StmtExprVisitor;
-
- public:
- explicit HoistInfoCollector(HoistExpressionConfig config) : config(config) {}
-
- private:
- void AttemptHoistConditional(PrimExpr cond, HoistedConditionals hoist_from,
- bool generate_else_block = true) {
- if (SideEffect(cond) > CallEffectKind::kPure) {
- return;
- }
- if (auto info = FindHoistDestination(cond)) {
- if (!info->reached_sequential_node) {
- // Record whether this conditional uses any block variables.
- auto walkfn = [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
- return active_block_vars.count(var.get())
- ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
- : ffi::WalkResult::Advance();
- };
- bool uses_block_var =
- active_block_vars.size() &&
- ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(cond,
walkfn).has_value();
-
- std::unordered_set<const VarNode*> let_bindings_used;
-
- for (Var var : UndefinedVars(cond)) {
- auto it = let_var_to_let_vars.find(var.get());
- if (it != let_var_to_let_vars.end()) {
- let_bindings_used.insert(it->first);
- for (auto used : it->second) {
- let_bindings_used.insert(used);
- }
- }
- }
- info->conditions.push_back(ConditionInfo(cond, hoist_from,
uses_block_var,
- let_bindings_used,
generate_else_block));
- }
- }
- }
-
- ffi::Optional<VisitInterrupt> Visit_(const prim::AndNode* op) final {
- AttemptHoistConditional(op->a, HoistedConditionals::kBooleanExpression);
- AttemptHoistConditional(op->b, HoistedConditionals::kBooleanExpression);
- return Parent::Visit_(op);
- }
-
- ffi::Optional<VisitInterrupt> Visit_(const prim::OrNode* op) final {
- AttemptHoistConditional(op->a, HoistedConditionals::kBooleanExpression);
- AttemptHoistConditional(op->b, HoistedConditionals::kBooleanExpression);
- return Parent::Visit_(op);
- }
-
- ffi::Optional<VisitInterrupt> Visit_(const ForNode* op) final {
- active_loops.push_back({op->loop_var, ffi::GetRef<Stmt>(op)});
- active_loop_vars.insert(op->loop_var.get());
-
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Parent::Visit_(op));
- completed_loops.push_back(active_loops.back());
-
- active_loop_vars.erase(op->loop_var.get());
- active_loops.pop_back();
- return std::nullopt;
- }
-
- ffi::Optional<VisitInterrupt> Visit_(const RegionStmtNode* op) final {
- if (op->op.same_as(tirx::launch_thread_op())) {
- Var var = op->body_params[0].as_or_throw<PrimVar>();
- active_block_vars.insert(var.get());
- active_loop_vars.insert(var.get());
- active_loops.push_back({var, ffi::GetRef<Stmt>(op)});
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Parent::Visit_(op));
- completed_loops.push_back(active_loops.back());
- active_loops.pop_back();
- active_loop_vars.erase(var.get());
- active_block_vars.erase(var.get());
- return std::nullopt;
- }
- // Unknown operations have no code-motion semantics. Analyze their inner
- // loops independently without hoisting across the region boundary.
- auto outer_loops = std::move(active_loops);
- auto outer_loop_vars = std::move(active_loop_vars);
- auto outer_block_vars = std::move(active_block_vars);
- active_loops.clear();
- active_loop_vars.clear();
- active_block_vars.clear();
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Parent::Visit_(op));
- active_loops = std::move(outer_loops);
- active_loop_vars = std::move(outer_loop_vars);
- active_block_vars = std::move(outer_block_vars);
- if (!active_loops.empty()) active_loops.back().reached_sequential_node =
true;
- return std::nullopt;
- }
-
- void VisitBinding(Var var, PrimExpr value, HoistedLetBindings hoist_from) {
- TVM_FFI_ICHECK_EQ(let_var_to_loop_vars.count(var.get()), 0)
- << "Multiple nested definitions of variable " << var;
- TVM_FFI_ICHECK_EQ(let_var_to_let_vars.count(var.get()), 0)
- << "Multiple nested definitions of variable " << var;
-
- if (auto info = FindHoistDestination(value)) {
- if (!info->reached_sequential_node) {
- info->let_bindings.push_back(LetBindingInfo(var, value, hoist_from));
- }
- }
-
- // Walk through the loop binding
- std::unordered_set<const VarNode*> loop_vars_used;
- std::unordered_set<const VarNode*> let_bindings_used;
- for (Var var : UndefinedVars(value)) {
- if (active_loop_vars.count(var.get())) {
- loop_vars_used.insert(var.get());
- } else {
- auto it = let_var_to_loop_vars.find(var.get());
- if (it != let_var_to_loop_vars.end()) {
- for (const VarNode* used : it->second) {
- loop_vars_used.insert(used);
- }
- }
- }
-
- auto it = let_var_to_let_vars.find(var.get());
- if (it != let_var_to_let_vars.end()) {
- let_bindings_used.insert(it->first);
- for (const VarNode* used : it->second) {
- let_bindings_used.insert(used);
- }
- }
- }
-
- let_var_to_loop_vars[var.get()] = std::move(loop_vars_used);
- let_var_to_let_vars[var.get()] = std::move(let_bindings_used);
- }
-
- ffi::Optional<VisitInterrupt> Visit_(const BindNode* op) final {
- if (auto value = op->value.as<PrimExpr>()) {
- VisitBinding(op->var, value.value(), HoistedLetBindings::kBind);
- }
- return Parent::Visit_(op);
- }
-
- ffi::Optional<VisitInterrupt> Visit_(const SeqStmtNode* op) final {
- if (active_loops.size()) {
- int non_bind_count = 0;
- for (size_t i = 0; i < op->seq.size(); ++i) {
- const auto* bind = op->seq[i].as<BindNode>();
- if (!bind) {
- non_bind_count++;
- } else if (const auto* call = bind->value.as<CallNode>();
- call && (call->op.same_as(tirx::alloc_tensor_op()) ||
- call->op.same_as(tirx::decl_tensor_op()))) {
- non_bind_count++;
- }
- }
- if (non_bind_count > 1) {
- active_loops.back().reached_sequential_node = true;
- }
- }
- std::vector<const VarNode*> seq_bind_vars;
- for (size_t i = 0; i < op->seq.size(); ++i) {
- if (auto* bind = op->seq[i].as<BindNode>()) {
- seq_bind_vars.push_back(bind->var.get());
- }
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(op->seq[i]));
- }
- for (auto* var : seq_bind_vars) {
- let_var_to_loop_vars.erase(var);
- let_var_to_let_vars.erase(var);
- }
- return std::nullopt;
- }
-
- ffi::Optional<VisitInterrupt> Visit_(const prim::LetNode* op) final {
- VisitBinding(op->var, op->value, HoistedLetBindings::kLetExpr);
-
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Parent::Visit_(op));
-
- let_var_to_loop_vars.erase(op->var.get());
- let_var_to_let_vars.erase(op->var.get());
- return std::nullopt;
- }
-
- ffi::Optional<VisitInterrupt> Visit_(const IfNode* op) final {
- AttemptHoistConditional(op->condition, HoistedConditionals::kIfElseStmt,
- op->else_case.has_value());
- return Parent::Visit_(op);
- }
-
- ffi::Optional<VisitInterrupt> Visit_(const CallNode* op) final {
- if (op->op.same_as(prim::if_then_else_op())) {
- PrimExpr cond = op->args[0].as_or_throw<PrimExpr>();
- AttemptHoistConditional(cond, HoistedConditionals::kIfElseExpr);
- }
- return Parent::Visit_(op);
- }
-
- // Find the loop above which this expression could be hoisted. If
- // nullptr, the expression cannot be hoisted.
- HoistInfo* FindHoistDestination(PrimExpr expr) {
- // Cannot hoist above a loop if we aren't already in a loop.
- if (active_loops.empty()) {
- return nullptr;
- }
-
- for (auto it = active_loops.rbegin(); it != active_loops.rend(); it++) {
- Var loop_var = it->loop_var;
- auto walkfn = [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
- bool matches = var.get() == loop_var.get();
- if (!matches) {
- auto it = let_var_to_loop_vars.find(var.get());
- matches = it != let_var_to_loop_vars.end() &&
it->second.count(loop_var.get());
- }
- return matches ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
- : ffi::WalkResult::Advance();
- };
- bool uses_loop_var =
ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(expr, walkfn).has_value();
-
- bool is_disabled_hoist_across_block_var =
- !config->FlagSet(HoistedConditionals::kUsingBlockVar) &&
it->IsBlockVariable();
-
- if (it->reached_sequential_node || uses_loop_var ||
is_disabled_hoist_across_block_var) {
- if (it == active_loops.rbegin()) {
- // Cannot hoist beyond the innermost loop iterator.
- return nullptr;
- } else {
- // Hoist to just below the loop iterator that is required.
- it--;
- return &(*it);
- }
- }
- }
-
- // If no loop variables are used, can hoist above the outermost
- // loop.
- return &active_loops.front();
- }
-
- // The user-provided config describing which expressions should be
- // hoisted.
- HoistExpressionConfig config;
-
- // Current hardware launch bindings of block variables.
- std::unordered_set<const VarNode*> active_block_vars;
-
- // An ordered list of loops that are currently being visited.
- std::vector<HoistInfo> active_loops;
-
- // Loops that have already been visited
- std::vector<HoistInfo> completed_loops;
-
- // Map from a bound variable to the loop variables it depends on.
- // Includes indirect usage.
- std::unordered_map<const VarNode*, std::unordered_set<const VarNode*>>
let_var_to_loop_vars;
-
- // Map from a bound variable to the other let bindings it depends on.
- // Includes indirect usage.
- std::unordered_map<const VarNode*, std::unordered_set<const VarNode*>>
let_var_to_let_vars;
-
- // Lookup table for the currently active loops.
- std::unordered_set<const VarNode*> active_loop_vars;
-};
-
-class ExpressionHoister : public s_tir::IRMutatorWithAnalyzer {
- public:
- using s_tir::IRMutatorWithAnalyzer::Mutate;
- using s_tir::IRMutatorWithAnalyzer::Mutate_;
-
- static Stmt Hoist(Stmt stmt, HoistExpressionConfig config) {
- auto loop_info = HoistInfoCollector::Collect(stmt, config);
-
- sym::Analyzer analyzer;
- auto hoister = ffi::make_object<ExpressionHoister>(std::move(loop_info),
config, analyzer);
- stmt = hoister->Mutate(stmt,
InplaceMode::kAllow).ValueOrUnchanged(std::move(stmt));
- stmt = s_tir::ConvertSSA(std::move(stmt));
- return stmt;
- }
-
- private:
- using Parent = s_tir::IRMutatorWithAnalyzer;
-
- public:
- explicit ExpressionHoister(std::vector<HoistInfoCollector::HoistInfo>
loop_info,
- HoistExpressionConfig config, const
sym::Analyzer& analyzer)
- : Parent(analyzer), config_(config) {
- for (auto& info : loop_info) {
- // Mark let bindings to use if they are enabled on their own.
- for (const auto& binding : info.let_bindings) {
- if (binding.IsEnabled(config)) {
- hoisted_let_bindings.insert(binding.var.get());
- }
- }
-
- // Or if they are required by a conditional
- if (config->FlagSet(HoistedLetBindings::kRequiredByCondition)) {
- for (const auto& conditional : info.conditions) {
- if (conditional.IsEnabled(config)) {
- for (const auto& var : conditional.required_let_bindings) {
- hoisted_let_bindings.insert(var);
- }
- }
- }
- }
-
- loop_info_lookup.insert_or_assign(info.loop_def.get(), std::move(info));
- }
- }
-
- private:
- Stmt WrapHoistedStatements(Stmt stmt, const HoistInfoCollector::HoistInfo&
info) {
- for (auto cond_it = info.conditions.rbegin(); cond_it !=
info.conditions.rend(); cond_it++) {
- if (cond_it->IsEnabled(config_)) {
- if (cond_it->generate_else_case) {
- stmt = If(cond_it->condition, stmt, SeqStmt(stmt));
- } else {
- stmt = If(cond_it->condition, stmt);
- }
- }
- }
- {
- ffi::Array<Stmt> binds;
- for (auto let_it = info.let_bindings.begin(); let_it !=
info.let_bindings.end(); let_it++) {
- if (hoisted_let_bindings.count(let_it->var.get())) {
- binds.push_back(Bind(let_it->var, let_it->value));
- }
- }
- if (!binds.empty()) {
- binds.push_back(stmt);
- stmt = SeqStmt(binds);
- }
- }
-
- return stmt;
- }
-
- UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final
{
- Stmt stmt = Parent::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
-
- auto it = loop_info_lookup.find(op);
- TVM_FFI_ICHECK(it != loop_info_lookup.end())
- << "Could not find pre-pass information for loop over " <<
op->loop_var;
- return WrapHoistedStatements(stmt, it->second);
- }
-
- UnchangedOr<Stmt> Mutate_(const RegionStmtNode* op, InplaceMode
inplace_mode) final {
- Stmt stmt = Parent::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
- auto it = loop_info_lookup.find(op);
- return it == loop_info_lookup.end() ? stmt : WrapHoistedStatements(stmt,
it->second);
- }
-
- UnchangedOr<Stmt> Mutate_(const BindNode* op, InplaceMode inplace_mode)
final {
- if (hoisted_let_bindings.count(op->var.get())) {
- // The binding was hoisted; remove it from this location.
- return Evaluate(0);
- } else {
- return Parent::Mutate_(op, inplace_mode);
- }
- }
-
- UnchangedOr<PrimExpr> Mutate_(const prim::LetNode* op, InplaceMode
inplace_mode) final {
- if (hoisted_let_bindings.count(op->var.get())) {
- return this->Mutate(op->body, inplace_mode).ValueOrUnchanged(op->body);
- } else {
- return Parent::Mutate_(op, inplace_mode);
- }
- }
-
- HoistExpressionConfig config_;
-
- std::unordered_map<const StmtNode*, HoistInfoCollector::HoistInfo>
loop_info_lookup;
- std::unordered_set<const VarNode*> hoisted_let_bindings;
-};
-
-Stmt HoistExpression(Stmt stmt, HoistExpressionConfig config) {
- return ExpressionHoister::Hoist(stmt, config);
-}
-
-namespace transform {
-
-Pass HoistExpression() {
- auto pass_func = [=](Function f, IRModule m, PassContext ctx) {
- if (!f->body.has_value()) return f;
- auto* n = f.CopyOnWrite();
- auto cfg =
ctx->GetConfigOrDefault<HoistExpressionConfig>("s_tir.HoistExpression");
- n->body = ExpressionHoister::Hoist(std::move(n->body).value(), cfg);
- return f;
- };
- auto insertion_pass = CreateFunctionPass(pass_func, 0,
"s_tir.InsertHoistedExpression");
-
- return tvm::transform::Sequential(
- {
- insertion_pass,
- s_tir::transform::StmtSimplify(),
- tirx::transform::RemoveNoOp(),
- },
- "s_tir.HoistExpression");
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("s_tir.transform.HoistExpression", HoistExpression);
-}
-
-static Pass HoistIfThenElseImpl() {
- auto pass_func = [=](Function f, IRModule m, PassContext ctx) {
- if (!f->body.has_value()) return f;
- auto* n = f.CopyOnWrite();
- auto cfg =
ctx->GetConfigOrDefault<HoistIfThenElseConfig>("s_tir.HoistIfThenElse");
- auto flag =
f->GetAttr<int64_t>(tvm::s_tir::attr::kHoistIfThenElseExprWithBlock);
- if (flag && flag.value() == 1) {
- HoistExpressionConfig
config(static_cast<int>(HoistedConditionals::kUsingBlockVar) |
-
static_cast<int>(HoistedConditionals::kIfElseExpr),
-
static_cast<int>(HoistedLetBindings::kNone));
- n->body = ExpressionHoister::Hoist(std::move(n->body).value(), config);
- return f;
- }
- int block_var =
- static_cast<int>(cfg->support_block_scope_hoisting ?
HoistedConditionals::kUsingBlockVar
- :
HoistedConditionals::kNone);
- HoistExpressionConfig config(block_var |
static_cast<int>(HoistedConditionals::kIfElseStmt),
- static_cast<int>(HoistedLetBindings::kNone));
- n->body = ExpressionHoister::Hoist(std::move(n->body).value(), config);
- return f;
- };
- auto insertion_pass = CreateFunctionPass(pass_func, 0,
"s_tir.InsertHoistIfThenElse");
- return tvm::transform::Sequential(
- {
- insertion_pass,
- s_tir::transform::StmtSimplify(),
- tirx::transform::RemoveNoOp(),
- },
- "s_tir.HoistIfThenElse");
-}
-
-static Pass HoistIfThenElseBasicImpl() {
- auto pass_func = [=](Function f, IRModule m, PassContext ctx) {
- if (!f->body.has_value()) return f;
- auto* n = f.CopyOnWrite();
- HoistExpressionConfig
config(static_cast<int>(HoistedConditionals::kIfElseStmt),
- static_cast<int>(HoistedLetBindings::kNone));
- n->body = ExpressionHoister::Hoist(std::move(n->body).value(), config);
- return f;
- };
- auto insertion_pass = CreateFunctionPass(pass_func, 0,
"s_tir.InsertHoistIfThenElseBasic");
- return tvm::transform::Sequential(
- {
- insertion_pass,
- s_tir::transform::StmtSimplify(),
- tirx::transform::RemoveNoOp(),
- },
- "s_tir.HoistIfThenElseBasic");
-}
-
-Pass HoistIfThenElse(tvm::ffi::String variant) {
- if (variant == "basic") {
- return HoistIfThenElseBasicImpl();
- }
- return HoistIfThenElseImpl();
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("s_tir.transform.HoistIfThenElse",
HoistIfThenElseImpl);
- refl::GlobalDef().def("s_tir.transform.HoistIfThenElseBasic",
HoistIfThenElseBasicImpl);
-}
-
-} // namespace transform
-
-} // namespace s_tir
-} // namespace tvm
diff --git a/src/tirx/transform/hoist_if.cc b/src/tirx/transform/hoist_if.cc
new file mode 100644
index 0000000000..09b01166f6
--- /dev/null
+++ b/src/tirx/transform/hoist_if.cc
@@ -0,0 +1,264 @@
+/*
+ * 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 hoist_if.cc
+ * \brief Hoist loop-invariant If statements.
+ */
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/prim/op.h>
+#include <tvm/ir/scope_stack.h>
+#include <tvm/tirx/analysis.h>
+#include <tvm/tirx/op/region.h>
+#include <tvm/tirx/stmt_functor.h>
+#include <tvm/tirx/transform.h>
+
+#include <algorithm>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include "ir_utils.h"
+
+namespace tvm {
+namespace tirx {
+
+class IfHoister : public StmtExprMutator {
+ public:
+ using StmtExprMutator::Mutate;
+ using StmtExprMutator::Mutate_;
+
+ UnchangedOr<ffi::Any> Mutate(ffi::AnyView input, InplaceMode inplace_mode)
final {
+ // Expressions are inspected only as whole If predicates, never split or
hoisted.
+ if (input.as<ExprNode>()) return ffi::Unchanged();
+ if (input.as<StmtNode>() && !input.as<ForNode>() && !input.as<IfNode>() &&
+ !input.as<SeqStmtNode>()) {
+ // Bindings, regions and unknown statements have no code-motion contract.
+ // Hide outer loops and their variable depths while visiting this
subtree;
+ // active_loops_/split_in_nest_ still bound duplication of the enclosing
nest.
+ return scopes_.WithNewScope([&] { return StmtExprMutator::Mutate(input,
inplace_mode); });
+ }
+ return StmtExprMutator::Mutate(input, inplace_mode);
+ }
+
+ private:
+ struct LoopState {
+ // Enclosing alternate-branch depth at loop entry, before visiting its
body.
+ int else_depth_at_entry;
+ // Predicates lifted to this loop, in traversal order; bool retains a
false-path copy.
+ std::vector<std::pair<PrimExpr, bool>> conditions;
+ };
+
+ struct ScopeState {
+ // Active destination loops, outermost first, within the current motion
boundary.
+ std::vector<LoopState*> loops;
+ // Identity-based indices into loops for variables defined by those loops.
+ std::unordered_map<const VarNode*, size_t> loop_depths;
+ };
+
+ UnchangedOr<Stmt> Mutate_(const SeqStmtNode* op, InplaceMode inplace_mode)
final {
+ // Sibling statements isolate loop placement; a singleton stays
transparent.
+ // An If inside one sibling must not move outside the loop executing all
siblings.
+ if (op->seq.size() == 1) return StmtExprMutator::Mutate_(op, inplace_mode);
+ return scopes_.WithNewScope([&] { return StmtExprMutator::Mutate_(op,
inplace_mode); });
+ }
+
+ // Purity alone does not permit evaluating a predicate before its guards or
+ // before a zero-trip loop. Check safety and loop dependencies in one walk.
+ size_t FindLiftDestination(const PrimExpr& condition) const {
+ const auto& scope = scopes_.Current();
+ // First loop the predicate can cross: loops [i, j, k] and i + j select
index 2 (k).
+ size_t destination = 0;
+ std::unordered_set<const ExprNode*> visited;
+ auto advance = [&](const ExprNode* op) -> ffi::Expected<ffi::WalkResult> {
+ // StructuralWalk visits occurrences. For Add(e, e), safety and
dependencies
+ // need checking only once per shared e, within this predicate's fixed
scope.
+ return visited.insert(op).second ? ffi::WalkResult::Advance() :
ffi::WalkResult::Skip();
+ };
+ auto division = [&](const ExprNode* op,
+ const PrimExpr& divisor) ->
ffi::Expected<ffi::WalkResult> {
+ // Positive constant divisors exclude both zero and signed min / -1.
+ const auto* value = divisor.as<IntImmNode>();
+ if (!value || value->value <= 0) return ffi::WalkResult::Interrupt();
+ return advance(op);
+ };
+ auto blocked = ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+ condition,
+ [&](const VarNode* op) -> ffi::Expected<ffi::WalkResult> {
+ // Stay inside every referenced loop variable's definition. Other
+ // variables remain in scope at every candidate destination.
+ auto it = scope.loop_depths.find(op);
+ if (it != scope.loop_depths.end()) destination =
std::max(destination, it->second + 1);
+ // An innermost-loop dependency already rules out every destination.
+ if (destination == scope.loops.size()) return
ffi::WalkResult::Interrupt();
+ // Variable types describe metadata, not evaluated predicate
dependencies.
+ return ffi::WalkResult::Skip();
+ },
+ [](const TypeNode*) -> ffi::Expected<ffi::WalkResult> { return
ffi::WalkResult::Skip(); },
+ [](const TensorLoadNode*) -> ffi::Expected<ffi::WalkResult> {
+ // Memory reads can change across iterations or depend on an
execution guard.
+ return ffi::WalkResult::Interrupt();
+ },
+ [&](const prim::DivNode* op) { return division(op, op->b); },
+ [&](const prim::ModNode* op) { return division(op, op->b); },
+ [&](const prim::FloorDivNode* op) { return division(op, op->b); },
+ [&](const prim::FloorModNode* op) { return division(op, op->b); },
+ [](const prim::LShiftNode*) -> ffi::Expected<ffi::WalkResult> {
+ // A shift count may be outside the defined range.
+ return ffi::WalkResult::Interrupt();
+ },
+ [](const prim::RShiftNode*) -> ffi::Expected<ffi::WalkResult> {
+ // A shift count may be outside the defined range.
+ return ffi::WalkResult::Interrupt();
+ },
+ [&](const prim::CastNode* op) -> ffi::Expected<ffi::WalkResult> {
+ // Floating-to-integer conversion may be undefined outside the
integer range.
+ if (ffi::GetRef<PrimExpr>(op).ty().MatchesCode(kDLInt, kDLUInt) &&
+ !op->value.ty().MatchesCode(kDLInt, kDLUInt))
+ return ffi::WalkResult::Interrupt();
+ return advance(op);
+ },
+ [&](const CallNode* op) -> ffi::Expected<ffi::WalkResult> {
+ // Even pure calls may trap; only the transparent likely hint is
known safe.
+ if (!op->op.same_as(prim::likely_op())) return
ffi::WalkResult::Interrupt();
+ return advance(op);
+ },
+ advance);
+ return blocked.has_value() ? scope.loops.size() : destination;
+ }
+
+ UnchangedOr<Stmt> Mutate_(const RegionStmtNode* op, InplaceMode
inplace_mode) final {
+ // Moving guards inside a thread launch can prevent later synchronization
+ // insertion, even when the guards themselves are loop-invariant.
+ if (op->op.same_as(launch_thread_op())) return ffi::Unchanged();
+ return StmtExprMutator::Mutate_(op, inplace_mode);
+ }
+
+ UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final
{
+ // Preserve the same synchronization boundary before thread-binding
lowering.
+ if (GetThreadBinding(op).has_value()) return ffi::Unchanged();
+ // Only ordinary serial loops permit code motion across their execution
scope.
+ // Hoisting must not suppress evaluation of an effectful loop header
either.
+ if (op->kind != ForKind::kDefault || SideEffect(op->min) >
CallEffectKind::kReadState ||
+ SideEffect(op->extent) > CallEffectKind::kReadState ||
+ (op->step && SideEffect(op->step.value()) >
CallEffectKind::kReadState)) {
+ return scopes_.WithNewScope([&] { return StmtExprMutator::Mutate_(op,
inplace_mode); });
+ }
+ auto& scope = scopes_.Current();
+ // A new outermost eligible loop gets one two-sided split. Do not reset at
+ // scope barriers: an outer i / SeqStmt / inner j still shares this budget.
+ if (active_loops_++ == 0) split_in_nest_ = false;
+ // Snapshot ancestors outside this loop so they do not require a
false-path copy.
+ LoopState loop{else_depth_, {}};
+ const VarNode* var = op->loop_var.get();
+ // Register this destination and its variable together. An index in loops
+ // identifies both where to wrap a lifted If and which variables it must
avoid.
+ scope.loop_depths[var] = scope.loops.size();
+ scope.loops.push_back(&loop);
+ Stmt stmt = StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
+ // Only descendants can lift to this loop; remove both entries before
siblings.
+ // The local LoopState remains alive below to emit its collected
conditions.
+ scope.loops.pop_back();
+ scope.loop_depths.erase(var);
+ // Reaching zero lets the next independent nest start with a fresh split
budget.
+ --active_loops_;
+ // Collected p then q must become if p: if q: loop, so wrap q before p.
+ for (auto it = loop.conditions.rbegin(); it != loop.conditions.rend();
++it) {
+ // Duplicate only when an alternate branch must remain reachable.
+ const auto& [condition, preserve_else] = *it;
+ stmt = preserve_else ? If(condition, stmt, SeqStmt(stmt)) :
If(condition, stmt);
+ }
+ return stmt;
+ }
+
+ UnchangedOr<Stmt> Mutate_(const IfNode* op, InplaceMode inplace_mode) final {
+ auto& scope = scopes_.Current();
+ // Moving a one-sided inner If could otherwise skip this predicate's
effects.
+ if (!scope.loops.empty() && SideEffect(op->condition) >
CallEffectKind::kReadState) {
+ return scopes_.WithNewScope([&] { return StmtExprMutator::Mutate_(op,
inplace_mode); });
+ }
+ bool has_else = op->else_case.has_value();
+ if (!scope.loops.empty()) {
+ size_t destination = FindLiftDestination(op->condition);
+ if (destination < scope.loops.size()) {
+ // A false predicate must still execute an enclosing alternate branch
+ // inside the destination loop, even when this If has no else of its
own.
+ // In for i: if p(i) { if q { A } } else { B }, q=false still needs B.
+ // The entry snapshot excludes Ifs outside this loop,
+ // whose alternate branches remain reachable without duplicating it.
+ bool preserve_else =
+ has_else || else_depth_ >
scope.loops[destination]->else_depth_at_entry;
+ auto* loop = scope.loops[destination];
+ // Each two-sided predicate doubles the enclosing loop subtree. Limit
+ // splitting to once per loop nest, even across code-motion barriers.
+ if (!preserve_else || !split_in_nest_) {
+ // Record at the destination now; its For unwind wraps the complete
+ // rewritten body. Preorder collection preserves enclosing predicate
order.
+ loop->conditions.emplace_back(op->condition, preserve_else);
+ // Hoists needing no false-path copy spend no budget. A copy consumes
+ // it for the rest of this nest, including loops behind scope
barriers.
+ split_in_nest_ |= preserve_else;
+ }
+ }
+ }
+ // Descendant hoists must preserve this If's alternate arm, in either
branch.
+ else_depth_ += has_else;
+ auto result = StmtExprMutator::Mutate_(op, inplace_mode);
+ // Restore the ancestor count so following siblings do not inherit this If.
+ else_depth_ -= has_else;
+ return result;
+ }
+
+ // Number of enclosing Ifs with an else; each loop saves its entry baseline.
+ int else_depth_{0};
+ // Active eligible loops across motion boundaries; zero starts a new split
budget.
+ // Unlike scoped destinations, this stays nonzero inside a nested SeqStmt
barrier.
+ int active_loops_{0};
+ // Whether this loop nest has used its one two-sided split, bounding subtree
growth.
+ bool split_in_nest_{false};
+ // Scoped destinations/dependencies; barriers hide outer loops without
resetting the budget.
+ ScopeStack<ScopeState> scopes_;
+};
+
+namespace transform {
+
+Pass HoistIf() {
+ auto pass_func = [](Function f, IRModule m, PassContext ctx) {
+ if (!f->body.has_value()) return f;
+ auto* n = f.CopyOnWrite();
+ Stmt body = ffi::make_object<IfHoister>()
+ ->Mutate(n->body.value(), InplaceMode::kAllow)
+ .ValueOrUnchanged(n->body.value());
+ n->body = tirx::ConvertSSA(std::move(body));
+ return f;
+ };
+ return tvm::transform::Sequential(
+ {CreateFunctionPass(pass_func, 0, "tirx.InsertHoistedIf"),
StmtSimplify(), RemoveNoOp()},
+ "tirx.HoistIf");
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
+ refl::GlobalDef().def("tirx.transform.HoistIf", HoistIf);
+}
+
+} // namespace transform
+} // namespace tirx
+} // namespace tvm
diff --git a/src/tirx/transform/remove_no_op.cc
b/src/tirx/transform/remove_no_op.cc
index 1a1790d21c..6e7419d48c 100644
--- a/src/tirx/transform/remove_no_op.cc
+++ b/src/tirx/transform/remove_no_op.cc
@@ -26,7 +26,6 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/op.h>
#include <tvm/ir/prim/op.h>
-#include <tvm/s_tir/stmt.h>
#include <tvm/sym/analyzer.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/stmt.h>
diff --git a/tests/python/s_tir/dlight/test_gpu_low_batch_gemv.py
b/tests/python/s_tir/dlight/test_gpu_low_batch_gemv.py
index 8cf6eebf66..ba2d320160 100644
--- a/tests/python/s_tir/dlight/test_gpu_low_batch_gemv.py
+++ b/tests/python/s_tir/dlight/test_gpu_low_batch_gemv.py
@@ -34,7 +34,7 @@ def test_batch_decode_gemv():
@Ts.function(private=True)
def before(lv429: T.Tensor((T.int64(4096), T.int64(3584)), "uint32"),
lv430: T.Tensor((T.int64(4096), T.int64(896)), "float16"), lv807:
T.Tensor((batch_size, T.int64(1), T.int64(28672)), 'float16'),
NT_matmul_intermediate: T.Tensor((batch_size, T.int64(1), T.int64(4096)),
'float16')):
- T.func_attr({"tirx.noalias": True,
"tirx.HoistIfThenElseExprWithBlock": 1})
+ T.func_attr({"tirx.noalias": True})
# with Ts.sblock("root"):
compute = Ts.sblock_alloc_buffer((T.int64(4096), T.int64(28672)),
"float16")
@@ -64,7 +64,7 @@ def test_batch_decode_gemv():
@Ts.function(private=True)
def expected(lv429: T.Tensor((T.int64(4096), T.int64(3584)), "uint32"),
lv430: T.Tensor((T.int64(4096), T.int64(896)), "float16"), lv807:
T.Tensor((batch_size, T.int64(1), T.int64(28672)), 'float16'),
NT_matmul_intermediate: T.Tensor((batch_size, T.int64(1), T.int64(4096)),
'float16')):
- T.func_attr({"tirx.HoistIfThenElseExprWithBlock": 1,
"tirx.is_scheduled": True, "tirx.noalias": True})
+ T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
# with Ts.sblock("root"):
dequantize_intermediate_intermediate_local =
Ts.sblock_alloc_buffer((T.int64(4096), T.int64(28672)), "float16",
scope="local")
@@ -162,7 +162,7 @@ def test_batch_gemv():
@Ts.function(private=True)
def before(A: T.Tensor((batch_size, T.int64(1), T.int64(K)), 'float16'),
B: T.Tensor((T.int64(N), T.int64(K)), "float16"), NT_matmul:
T.Tensor((batch_size, T.int64(1), T.int64(N)), 'float16')):
- T.func_attr({"tirx.noalias": True,
"tirx.HoistIfThenElseExprWithBlock": 1})
+ T.func_attr({"tirx.noalias": True})
# with Ts.sblock("root"):
for i0, i1, i2, k in T.grid(batch_size, T.int64(1), T.int64(N),
T.int64(K)):
@@ -178,7 +178,7 @@ def test_batch_gemv():
@Ts.function(private=True)
def expected(A: T.Tensor((batch_size, T.int64(1), T.int64(4096)),
'float16'), B: T.Tensor((T.int64(4096), T.int64(4096)), "float16"), NT_matmul:
T.Tensor((batch_size, T.int64(1), T.int64(4096)), 'float16')):
- T.func_attr({"tirx.HoistIfThenElseExprWithBlock": 1,
"tirx.is_scheduled": True, "tirx.noalias": True})
+ T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
# with Ts.sblock("root"):
NT_matmul_pad_local = Ts.sblock_alloc_buffer(((batch_size +
T.int64(3)) // T.int64(4) * T.int64(4), T.int64(1), T.int64(4096)), "float16",
scope="local")
diff --git
a/tests/python/s_tir/transform/test_s_tir_transform_hoist_expression.py
b/tests/python/s_tir/transform/test_s_tir_transform_hoist_expression.py
deleted file mode 100644
index 4de0d34327..0000000000
--- a/tests/python/s_tir/transform/test_s_tir_transform_hoist_expression.py
+++ /dev/null
@@ -1,590 +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.
-# ruff: noqa: F401
-import enum
-
-import tvm
-import tvm.testing
-from tvm import s_tir
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
-
-
-class HoistedConditionals(enum.Flag):
- """Test flags for the native hoisted_conditionals pass option
-
- Each bitflag represents a type of expression that should be
- hoisted to the outermost loop possible.
- """
-
- Never = 0
- """ No hoisting of conditionals """
-
- IfElseStmt = 1
- """ If set, look for hoist candidates in IfElseStmt """
-
- IfElseExpr = 2
- """ If set, look for hoist candidates in tirx.if_then_else """
-
- BooleanExpression = 4
- """ If set, look for hoist candidates in all boolean expressions """
-
- UsingBlockVar = 8
- """ If set, allow hoisting of conditionals that use a block variable (e.g.
threadIdx.x) """
-
- All = IfElseStmt | IfElseExpr | BooleanExpression | UsingBlockVar
- """ Enable all hoisting of conditionals"""
-
-
-class HoistedLetBindings(enum.Flag):
- """Test flags for the native hoisted_let_bindings pass option
-
- Each bitflag represents a type of let binding expression that should be
- hoisted to the outermost loop possible.
- """
-
- Never = 0
- """ No hoisting of let bindings """
-
- RequiredByConditional = 1
- """ Bindings that are used by a hoisted conditional """
-
- Bind = 2
- """ Bindings occurring in Bind nodes """
-
- LetExpr = 4
- """ Bindings occurring in Let expressions """
-
- All = RequiredByConditional | Bind | LetExpr
- """ Enable all hoisting of let bindings """
-
-
-def _run_transform(before, hoisted_conditionals, hoisted_let_bindings):
- """Run HoistExpression transform with the given configuration."""
- before_mod = tvm.IRModule.from_expr(before)
-
- config = {
- "s_tir.HoistExpression": {
- "hoisted_conditionals": hoisted_conditionals.value,
- "hoisted_let_bindings": hoisted_let_bindings.value,
- }
- }
-
- with tvm.transform.PassContext(config=config):
- after_mod = tvm.s_tir.transform.HoistExpression()(before_mod)
-
- return after_mod["main"]
-
-
-def test_hoist_to_top_if_else_stmt():
- @Ts.function(private=True)
- def before(A: T.Tensor((16,), "float32"), n: T.int32):
- for i in T.serial(16):
- if n != 0:
- A[i] = 0.0
-
- @Ts.function(private=True)
- def expected(A: T.Tensor((16,), "float32"), n: T.int32):
- if n != 0:
- for i in T.serial(16):
- A[i] = 0.0
-
- after = _run_transform(before, HoistedConditionals.IfElseStmt,
HoistedLetBindings.All)
- tvm.ir.assert_structural_equal(after, expected)
-
-
-def test_hoist_to_top_all():
- @Ts.function(private=True)
- def before(A: T.Tensor((16,), "float32"), n: T.int32):
- for i in T.serial(16):
- if n != 0:
- A[i] = 0.0
-
- @Ts.function(private=True)
- def expected(A: T.Tensor((16,), "float32"), n: T.int32):
- if n != 0:
- for i in T.serial(16):
- A[i] = 0.0
-
- after = _run_transform(before, HoistedConditionals.All,
HoistedLetBindings.All)
- tvm.ir.assert_structural_equal(after, expected)
-
-
-def test_suppress_hoist_if_else_never():
- @Ts.function(private=True)
- def before(A: T.Tensor((16,), "float32"), n: T.int32):
- for i in T.serial(16):
- if n != 0:
- A[i] = 0.0
-
- expected = before
-
- after = _run_transform(before, HoistedConditionals.Never,
HoistedLetBindings.All)
- tvm.ir.assert_structural_equal(after, expected)
-
-
-def test_suppress_hoist_if_else_expr_only():
- @Ts.function(private=True)
- def before(A: T.Tensor((16,), "float32"), n: T.int32):
- for i in T.serial(16):
- if n != 0:
- A[i] = 0.0
-
- expected = before
-
- after = _run_transform(before, HoistedConditionals.IfElseExpr,
HoistedLetBindings.All)
- tvm.ir.assert_structural_equal(after, expected)
-
-
-def test_hoist_block_var():
- @Ts.function(private=True)
- def before(A: T.Tensor((128, 16), "float32"), n: T.int32):
- i = T.launch_thread("threadIdx.x", 128)
-
- for j in T.serial(16):
- if i < 32:
- A[i, j] = 0.0
-
- @Ts.function(private=True)
- def expected(A: T.Tensor((128, 16), "float32"), n: T.int32):
- i = T.launch_thread("threadIdx.x", 128)
-
- if i < 32:
- for j in T.serial(16):
- A[i, j] = 0.0
-
- after = _run_transform(before, HoistedConditionals.All,
HoistedLetBindings.All)
- tvm.ir.assert_structural_equal(after, expected)
-
-
-def test_suppress_hoist_block_var():
- @Ts.function(private=True)
- def before(A: T.Tensor((128, 16), "float32"), n: T.int32):
- T.launch_thread("threadIdx.x", 128)
-
- for i in T.thread_binding(0, 128, thread="threadIdx.x"):
- if i < 32:
- for j in T.serial(16):
- A[i, j] = 0.0
-
- expected = before
-
- after = _run_transform(
- before,
- HoistedConditionals.All & ~HoistedConditionals.UsingBlockVar,
- HoistedLetBindings.All,
- )
- tvm.ir.assert_structural_equal(after, expected)
-
-
-def test_hoist_across_block_var():
- @Ts.function(private=True)
- def before(A: T.Tensor((128, 16), "float32"), n: T.int32):
- T.launch_thread("threadIdx.x", 128)
-
- for i in T.thread_binding(0, 128, thread="threadIdx.x"):
- if n == 0:
- for j in T.serial(16):
- A[i, j] = 0.0
-
- @Ts.function(private=True)
- def expected(A: T.Tensor((128, 16), "float32"), n: T.int32):
- if n == 0:
- T.launch_thread("threadIdx.x", 128)
- for i in T.thread_binding(0, 128, thread="threadIdx.x"):
- for j in T.serial(16):
- A[i, j] = 0.0
-
- after = _run_transform(before, HoistedConditionals.All,
HoistedLetBindings.All)
- tvm.ir.assert_structural_equal(after, expected)
-
-
-def test_suppress_hoist_across_block_var():
- @Ts.function(private=True)
- def before(A: T.Tensor((128, 16), "float32"), n: T.int32):
- T.launch_thread("threadIdx.x", 128)
-
- for i in T.thread_binding(0, 128, thread="threadIdx.x"):
- for j in T.serial(16):
- if n == 0:
- A[i, j] = 0.0
-
- @Ts.function(private=True)
- def expected(A: T.Tensor((128, 16), "float32"), n: T.int32):
- T.launch_thread("threadIdx.x", 128)
- if n == 0:
- for i in T.thread_binding(0, 128, thread="threadIdx.x"):
- for j in T.serial(16):
- A[i, j] = 0.0
-
- after = _run_transform(
- before,
- HoistedConditionals.All & ~HoistedConditionals.UsingBlockVar,
- HoistedLetBindings.All,
- )
- tvm.ir.assert_structural_equal(after, expected)
-
-
-def test_hoist_to_middle():
- @Ts.function(private=True)
- def before(A: T.Tensor((4, 4), "float32")):
- for i in T.serial(4):
- for j in T.serial(4):
- if i < 3:
- A[i, j] = 0.0
-
- @Ts.function(private=True)
- def expected(A: T.Tensor((4, 4), "float32")):
- for i in T.serial(4):
- if i < 3:
- for j in T.serial(4):
- A[i, j] = 0.0
-
- after = _run_transform(before, HoistedConditionals.All,
HoistedLetBindings.All)
- tvm.ir.assert_structural_equal(after, expected)
-
-
-def test_hoist_with_let():
- @Ts.function(private=True)
- def before(A: T.Tensor((4, 4), "float32")):
- for i in T.serial(4):
- for j in T.serial(4):
- condition: T.let[T.bool] = i < 3
- if condition:
- A[i, j] = 0.0
-
- @Ts.function(private=True)
- def expected(A: T.Tensor((4, 4), "float32")):
- for i in T.serial(4):
- condition: T.let[T.bool] = i < 3 # noqa: F841
- if i < 3:
- for j in T.serial(4):
- A[i, j] = T.float32(0.0)
-
- after = _run_transform(before, HoistedConditionals.All,
HoistedLetBindings.All)
- tvm.ir.assert_structural_equal(after, expected)
-
-
-def test_hoist_disable_let():
- """As test_hoist_with_let, but forbid hoisting of Bind
-
- Because the condition depends on the let binding, it should no
- longer be hoisted. With Bind lifecycle management, the condition
- var is erased at sequence boundaries, so the if-condition uses
- the raw expression.
- """
-
- @Ts.function(private=True)
- def before(A: T.Tensor((4, 4), "float32")):
- for i in T.serial(4):
- for j in T.serial(4):
- condition: T.let[T.bool] = i < 3
- if condition:
- A[i, j] = 0.0
-
- @Ts.function(private=True)
- def expected(A: T.Tensor((4, 4), "float32")):
- for i, j in T.grid(4, 4):
- condition: T.let[T.bool] = i < 3 # noqa: F841
- if i < 3:
- A[i, j] = T.float32(0.0)
-
- after = _run_transform(before, HoistedConditionals.All,
HoistedLetBindings.Never)
- tvm.ir.assert_structural_equal(after, expected)
-
-
-def test_hoist_if_else():
- @Ts.function(private=True)
- def before(A: T.Tensor((4, 4), "float32")):
- for i in T.serial(4):
- for j in T.serial(4):
- if i < 3:
- A[i, j] = 0.0
- else:
- A[i, j] = 1.0
-
- @Ts.function(private=True)
- def expected(A: T.Tensor((4, 4), "float32")):
- for i in T.serial(4):
- if i < 3:
- for j in T.serial(4):
- A[i, j] = 0.0
- else:
- for j in T.serial(4):
- A[i, j] = 1.0
-
- after = _run_transform(before, HoistedConditionals.All,
HoistedLetBindings.All)
- tvm.ir.assert_structural_equal(after, expected)
-
-
-def test_hoist_sequential_assign():
- @Ts.function(private=True)
- def before(A: T.Tensor((4, 4), "float32"), B: T.Tensor((4, 4), "float32")):
- for i in T.serial(4):
- for j in T.serial(4):
- if i < 3:
- A[i, j] = 0.0
- B[i, j] = 0.0
- else:
- A[i, j] = 1.0
- B[i, j] = 1.0
-
- @Ts.function(private=True)
- def expected(A: T.Tensor((4, 4), "float32"), B: T.Tensor((4, 4),
"float32")):
- for i in T.serial(4):
- if i < 3:
- for j in T.serial(4):
- A[i, j] = 0.0
- B[i, j] = 0.0
- else:
- for j in T.serial(4):
- A[i, j] = 1.0
- B[i, j] = 1.0
-
- after = _run_transform(before, HoistedConditionals.All,
HoistedLetBindings.All)
- tvm.ir.assert_structural_equal(after, expected)
-
-
-def test_hoist_multi_if():
- @Ts.function(private=True)
- def before(A: T.Tensor((4, 4), "float32")):
- for i in T.serial(4):
- for j in T.serial(4):
- for k in T.serial(4):
- if j < 3:
- if i < 2:
- A[i, j] = 0.0
-
- @Ts.function(private=True)
- def expected(A: T.Tensor((4, 4), "float32")):
- for i in T.serial(4):
- if i < 2:
- for j in T.serial(4):
- if j < 3:
- for k in T.serial(4):
- A[i, j] = 0.0
-
- after = _run_transform(before, HoistedConditionals.All,
HoistedLetBindings.All)
- tvm.ir.assert_structural_equal(after, expected)
-
-
-def test_hoist_complex_conditional():
- @Ts.function(private=True)
- def before(A: T.Tensor((4, 4), "float32")):
- for i, j, k in T.grid(4, 4, 4):
- if j < 3 and i < 2:
- A[i, j] = 0.0
-
- @Ts.function(private=True)
- def expected(A: T.Tensor((4, 4), "float32")):
- for i in T.serial(4):
- if i < 2:
- for j in T.serial(4):
- if j < 3:
- for k in T.serial(4):
- A[i, j] = 0.0
-
- after = _run_transform(before, HoistedConditionals.All,
HoistedLetBindings.All)
- tvm.ir.assert_structural_equal(after, expected)
-
-
-def test_suppress_splitting_conditional():
- @Ts.function(private=True)
- def before(A: T.Tensor((4, 4), "float32")):
- for i, j, k in T.grid(4, 4, 4):
- if j < 3 and i < 2:
- A[i, j] = 0.0
-
- @Ts.function(private=True)
- def expected(A: T.Tensor((4, 4), "float32")):
- for i, j in T.grid(4, 4):
- if j < 3 and i < 2:
- for k in T.serial(4):
- A[i, j] = 0.0
-
- after = _run_transform(
- before,
- HoistedConditionals.All & ~HoistedConditionals.BooleanExpression,
- HoistedLetBindings.All,
- )
- tvm.ir.assert_structural_equal(after, expected)
-
-
-def test_hoist_multi_if_else():
- @Ts.function(private=True)
- def before(A: T.Tensor((4, 4), "float32")):
- for i in T.serial(4):
- for j in T.serial(4):
- for k in T.serial(4):
- if j < 3:
- if i < 2:
- A[i, j] = 0.0
- else:
- A[i, j] = 1.0
- else:
- if i < 2:
- A[i, j] = 2.0
- else:
- A[i, j] = 3.0
-
- @Ts.function(private=True)
- def expected(A: T.Tensor((4, 4), "float32")):
- for i in T.serial(4):
- if i < 2:
- for j in T.serial(4):
- if j < 3:
- for k in T.serial(4):
- A[i, j] = 0.0
- else:
- for k in T.serial(4):
- A[i, j] = 2.0
- else:
- for j in T.serial(4):
- if j < 3:
- for k in T.serial(4):
- A[i, j] = 1.0
- else:
- for k in T.serial(4):
- A[i, j] = 3.0
-
- after = _run_transform(before, HoistedConditionals.All,
HoistedLetBindings.All)
- tvm.ir.assert_structural_equal(after, expected)
-
-
-def test_hoist_multi_if_else_different_branches():
- @Ts.function(private=True)
- def before(A: T.Tensor((4, 4), "float32")):
- for i in T.serial(4):
- for j in T.serial(4):
- for k in T.serial(4):
- if j < 3:
- if i < 2:
- A[i, j] = 0.0
- else:
- A[i, j] = 1.0
- else:
- if i < 1:
- A[i, j] = 2.0
- else:
- A[i, j] = 3.0
-
- @Ts.function(private=True)
- def expected(A: T.Tensor((4, 4), "float32")):
- for i in T.serial(4):
- if i < 2:
- if i < 1:
- for j in T.serial(4):
- if j < 3:
- for k in T.serial(4):
- A[i, j] = 0.0
- else:
- for k in T.serial(4):
- A[i, j] = 2.0
- else:
- for j in T.serial(4):
- if j < 3:
- for k in T.serial(4):
- A[i, j] = 0.0
- else:
- for k in T.serial(4):
- A[i, j] = 3.0
- else:
- for j in T.serial(4):
- if j < 3:
- for k in T.serial(4):
- A[i, j] = 1.0
- else:
- for k in T.serial(4):
- A[i, j] = 3.0
-
- after = _run_transform(before, HoistedConditionals.All,
HoistedLetBindings.All)
- tvm.ir.assert_structural_equal(after, expected)
-
-
-def test_hoist_if_else_expr():
- @Ts.function(private=True)
- def before(A: T.Tensor((4, 4), "float32")):
- for i, j in T.grid(4, 4):
- A[i, j] = T.if_then_else(i < 2, 1.0, 2.0)
-
- @Ts.function(private=True)
- def expected(A: T.Tensor((4, 4), "float32")):
- for i in T.serial(4):
- if i < 2:
- for j in T.serial(4):
- A[i, j] = 1.0
- else:
- for j in T.serial(4):
- A[i, j] = 2.0
-
- after = _run_transform(before, HoistedConditionals.All,
HoistedLetBindings.All)
- tvm.ir.assert_structural_equal(after, expected)
-
-
-def test_suppress_hoist_if_else_expr():
- @Ts.function(private=True)
- def before(A: T.Tensor((4, 4), "float32")):
- for i, j in T.grid(4, 4):
- A[i, j] = T.if_then_else(i < 2, 1.0, 2.0)
-
- expected = before
-
- after = _run_transform(
- before,
- HoistedConditionals.All & ~HoistedConditionals.IfElseExpr,
- HoistedLetBindings.All,
- )
- tvm.ir.assert_structural_equal(after, expected)
-
-
-def test_hoist_let_expr():
- @Ts.function(private=True)
- def before(A: T.Tensor((4, 4), "float32")):
- for i, j in T.grid(4, 4):
- x = T.float32()
- A[i, j] = T.Let(5.0 * x + T.cast(j, "float32"), where={x: T.cast(i
+ 1, "float32")})
-
- @Ts.function(private=True)
- def expected(A: T.Tensor((4, 4), "float32")):
- for i in T.serial(4):
- x: T.let[T.float32] = T.cast(i + 1, "float32") # noqa: F841
- for j in T.serial(4):
- A[i, j] = T.float32(5.0) * T.cast(i + 1, "float32") +
T.cast(j, "float32")
-
- after = _run_transform(before, HoistedConditionals.All,
HoistedLetBindings.All)
- tvm.ir.assert_structural_equal(after, expected)
-
-
-def test_suppress_hoist_let_expr():
- @Ts.function(private=True)
- def before(A: T.Tensor((4, 4), "float32")):
- for i, j in T.grid(4, 4):
- x = T.float32()
- A[i, j] = T.Let(5.0 * x + T.cast(j, "float32"), where={x: T.cast(i
+ 1, "float32")})
-
- expected = before
-
- after = _run_transform(
- before,
- HoistedConditionals.All,
- HoistedLetBindings.All & ~HoistedLetBindings.LetExpr,
- )
- tvm.ir.assert_structural_equal(after, expected)
-
-
-if __name__ == "__main__":
- tvm.testing.main()
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
deleted file mode 100644
index 10116492b9..0000000000
--- a/tests/python/s_tir/transform/test_s_tir_transform_hoist_if.py
+++ /dev/null
@@ -1,562 +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.
-# ruff: noqa: E741, F401
-import numpy as np
-import pytest
-import tvm_ffi
-
-import tvm
-from tvm import s_tir
-from tvm.script import ir as I
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
-from tvm.testing import enabled_targets
-
-var_list = []
-
-
-def verify_structure(stmt, expected_struct):
- node_dict = {}
- struct = {}
-
- def _extract_vars(op):
- global var_list
- if isinstance(op, tvm.tirx.Var):
- var_list.append(op.name)
-
- def _visit(op):
- key = op
- if isinstance(op, tvm.ir.If):
- global var_list
- tvm_ffi.structural_walk(op.condition, _extract_vars)
- val = [(op.then_case, op.else_case), ("ir.If", tuple(var_list))]
- var_list.clear()
- elif isinstance(op, tvm.ir.For):
- val = [(op.body,), ("ir.For", op.loop_var.name)]
- elif isinstance(op, tvm.ir.RegionStmt):
- assert op.op.name == "tirx.launch_thread"
- val = [(op.body,), ("ir.RegionStmt", op.op.name, int(op.args[1]))]
- else:
- return
- node_dict[key] = val
-
- tvm_ffi.structural_walk(stmt, _visit)
- for key, val in node_dict.items():
- struct[val[1]] = tuple(
- node_dict[child[0]][1]
- if child is not None and len(child) == 1 and child[0] in node_dict
- else None
- for child in val[0]
- )
-
- assert struct == expected_struct, (
- f"Structure mismatch: expect {expected_struct} but got {struct}"
- )
- var_list.clear()
-
-
-def _opaque_eval(var):
- return tvm.ir.Evaluate(tvm.tirx.call_extern("int32", "dummy", var))
-
-
-def test_hoist_top_for():
- @Ts.function(private=True)
- def func(l: T.int32, m: T.int32, n: T.int32):
- for i in T.serial(l):
- for j in T.serial(m):
- for k in T.serial(n):
- if T.likely(i < 2):
- T.evaluate(T.call_extern("dummy", m, ty="int32"))
- else:
- T.evaluate(T.call_extern("dummy", n, ty="int32"))
-
- mod = tvm.IRModule.from_expr(func)
- new_stmt = tvm.s_tir.transform.HoistIfThenElse()(mod)["main"].body
- expected_struct = {
- ("ir.For", "k"): (None,),
- ("ir.For", "j"): (("ir.For", "k"),),
- ("ir.If", ("i",)): (("ir.For", "j"), ("ir.For", "j")),
- ("ir.For", "i"): (("ir.If", ("i",)),),
- }
- verify_structure(new_stmt, expected_struct)
-
-
-def test_hoist_multi_var_if():
- @Ts.function(private=True)
- def func(l: T.int32, m: T.int32, n: T.int32):
- for i in T.serial(l):
- for j in T.serial(m):
- for k in T.serial(n):
- if T.likely(i + j < 2):
- T.evaluate(T.call_extern("dummy", m, ty="int32"))
- else:
- T.evaluate(T.call_extern("dummy", n, ty="int32"))
-
- mod = tvm.IRModule.from_expr(func)
- new_mod = tvm.s_tir.transform.HoistIfThenElse()(mod)
- new_stmt = new_mod["main"].body
- expected_struct = {
- ("ir.For", "k"): (None,),
- ("ir.If", ("i", "j")): (("ir.For", "k"), ("ir.For", "k")),
- ("ir.For", "j"): (("ir.If", ("i", "j")),),
- ("ir.For", "i"): (("ir.For", "j"),),
- }
- verify_structure(new_stmt, expected_struct)
-
-
-def test_hoist_no_match_for():
- @Ts.function(private=True)
- def func(data: T.handle("float32"), l: T.int32, m: T.int32, n: T.int32):
- data_ptr = T.decl_tensor(1, "float32", data=data)
- for i in T.serial(l):
- for j in T.serial(m):
- data_ptr[i * 3 + j] = data_ptr[i * 3 + j] + T.float32(0.5)
- for k in T.serial(n):
- if T.likely(i < 2):
- T.evaluate(T.call_extern("dummy", m, ty="int32"))
- else:
- T.evaluate(T.call_extern("dummy", n, ty="int32"))
-
- mod = tvm.IRModule.from_expr(func)
- new_stmt = tvm.s_tir.transform.HoistIfThenElse()(mod)["main"].body
- expected_struct = {
- ("ir.For", "k"): (None,),
- ("ir.If", ("i",)): (("ir.For", "k"), ("ir.For", "k")),
- ("ir.For", "j"): (None,),
- ("ir.For", "i"): (("ir.For", "j"),),
- }
- verify_structure(new_stmt, expected_struct)
-
-
-def test_no_else():
- @Ts.function(private=True)
- def func(l: T.int32, m: T.int32, n: T.int32):
- for i in T.serial(l):
- for j in T.serial(m):
- for k in T.serial(n):
- if T.likely(i < 2):
- T.evaluate(T.call_extern("dummy", m, ty="int32"))
-
- mod = tvm.IRModule.from_expr(func)
- new_stmt = tvm.s_tir.transform.HoistIfThenElse()(mod)["main"].body
- expected_struct = {
- ("ir.For", "k"): (None,),
- ("ir.For", "j"): (("ir.For", "k"),),
- ("ir.If", ("i",)): (("ir.For", "j"), None),
- ("ir.For", "i"): (("ir.If", ("i",)),),
- }
- verify_structure(new_stmt, expected_struct)
-
-
-def test_thread_launch():
- dshape = (32, 64)
-
- @Ts.function(private=True)
- def func(data: T.handle("float32"), l: T.int32, m: T.int32, n: T.int32):
- data_ptr = T.decl_tensor(1, "float32", data=data)
- tx = T.launch_thread("threadIdx.x", dshape[0])
- bx = T.launch_thread("blockIdx.x", dshape[1])
- for i in T.serial(l):
- for j in T.serial(m):
- for k in T.serial(n):
- if i < 4 or j >= 8:
- data_ptr[bx * j + tx * j * k] = data_ptr[bx * j + tx *
j * k] + T.float32(
- 0.5
- )
- else:
- data_ptr[bx * j + tx * j * k] = data_ptr[bx * j + tx *
j * k] + T.float32(
- 1.0
- )
-
- mod = tvm.IRModule.from_expr(func)
- new_stmt = tvm.s_tir.transform.HoistIfThenElse()(mod)["main"].body
- expected_struct = {
- ("ir.For", "k"): (None,),
- ("ir.If", ("i", "j")): (("ir.For", "k"), ("ir.For", "k")),
- ("ir.For", "j"): (("ir.If", ("i", "j")),),
- ("ir.For", "i"): (("ir.For", "j"),),
- ("ir.RegionStmt", "tirx.launch_thread", 64): (("ir.For", "i"),),
- ("ir.RegionStmt", "tirx.launch_thread", 32): (("ir.RegionStmt",
"tirx.launch_thread", 64),),
- }
- verify_structure(new_stmt, expected_struct)
-
-
-def test_nested_for():
- @Ts.function(private=True)
- def func(data: T.handle("float32")):
- data_ptr = T.decl_tensor(1, "float32", data=data)
- for i in range(5):
- for j in range(10):
- if i >= 3:
- data_ptr[i * 3 + j] = data_ptr[i * 3 + j] + T.float32(0.5)
- for k in range(15):
- for l in range(20):
- if i < 4 or j >= 8:
- data_ptr[i * 3 + j + k + l] = data_ptr[
- i * 3 + j + k + l
- ] * T.float32(2)
- else:
- data_ptr[i * 3 + j + k + l] = data_ptr[
- i * 3 + j + k + l
- ] * T.float32(1.5)
-
- mod = tvm.IRModule.from_expr(func)
- new_stmt = tvm.s_tir.transform.HoistIfThenElse()(mod)["main"].body
- expected_struct = {
- ("ir.For", "l"): (None,),
- ("ir.For", "k"): (("ir.For", "l"),),
- ("ir.If", ("i", "j")): (("ir.For", "k"), ("ir.For", "k")),
- ("ir.For", "j"): (None,),
- ("ir.If", ("i",)): (("ir.For", "j"), None),
- ("ir.For", "i"): (("ir.If", ("i",)),),
- }
- verify_structure(new_stmt, expected_struct)
-
-
-def test_if_block():
- # Use different variable names for second loop nest to avoid dict key
collision
- @I.ir_module
- class Module:
- @Ts.function(private=True)
- def main(data: T.Tensor((1,), "float32"), n: T.int32):
- # First loop nest: i, j, k, l
- for i in T.serial(5):
- for j in T.serial(10):
- if i >= 3:
- data[i * 3 + j] = data[i * 3 + j] + T.float32(0.5)
- for k in T.serial(15):
- for l in T.serial(20):
- if i < 4 or j >= 8:
- data[i * 3 + j + k + l] = data[i * 3 + j +
k + l] * T.float32(2)
- else:
- data[i * 3 + j + k + l] = data[i * 3 + j +
k + l] * T.float32(
- 1.5
- )
- if j < 5:
- data[i * 3 + j + k + l] = data[i * 3 + j +
k + l] - T.float32(1)
-
- # Second loop nest: i2, j2, k2 (different names)
- for i2 in T.serial(5):
- for j2 in T.serial(10):
- for k2 in T.serial(15):
- if n >= 3:
- data[i2 * 3 + j2 + k2] = data[i2 * 3 + j2 + k2] +
T.float32(0.6)
-
- new_stmt = tvm.s_tir.transform.HoistIfThenElse()(Module)["main"].body
- # Updated expected_struct with renamed second nest variables
- expected_struct = {
- ("ir.If", ("i", "j")): (None, None),
- ("ir.If", ("j",)): (None, None),
- ("ir.For", "l"): (None,),
- ("ir.For", "k"): (("ir.For", "l"),),
- ("ir.For", "j"): (None,),
- ("ir.If", ("i",)): (("ir.For", "j"), None),
- ("ir.For", "i"): (("ir.If", ("i",)),),
- ("ir.For", "k2"): (None,),
- ("ir.For", "j2"): (("ir.For", "k2"),),
- ("ir.For", "i2"): (("ir.For", "j2"),),
- ("ir.If", ("n",)): (("ir.For", "i2"), None),
- }
- verify_structure(new_stmt, expected_struct)
-
-
-def test_multi_if():
- @Ts.function(private=True)
- def func(data: T.handle("float32")):
- data_ptr = T.decl_tensor(1, "float32", data=data)
- for i in range(10):
- for j in range(10):
- for k in range(10):
- if 3 <= i:
- if 3 <= j:
- data_ptr[i * 100 + j * 10 + k] = data_ptr[
- i * 100 + j * 10 + k
- ] + T.float32(0.5)
-
- mod = tvm.IRModule.from_expr(func)
- new_mod = tvm.s_tir.transform.HoistIfThenElse()(mod)
- new_stmt = new_mod["main"].body
- expected_struct = {
- ("ir.For", "k"): (None,),
- ("ir.If", ("j",)): (("ir.For", "k"), None),
- ("ir.For", "j"): (("ir.If", ("j",)),),
- ("ir.If", ("i",)): (("ir.For", "j"), None),
- ("ir.For", "i"): (("ir.If", ("i",)),),
- }
- verify_structure(new_stmt, expected_struct)
-
-
-def test_no_hoisting_1():
- @Ts.function(private=True)
- def func(data: T.handle("float32")):
- data_ptr = T.decl_tensor(1, "float32", data=data)
- for i in range(10):
- for j in range(10):
- for k in range(10):
- if k <= 3:
- data_ptr[i * 100 + j * 10 + k] = data_ptr[i * 100 + j
* 10 + k] + T.float32(
- 0.5
- )
-
- mod = tvm.IRModule.from_expr(func)
- stmt = mod["main"].body
- new_stmt = tvm.s_tir.transform.HoistIfThenElse()(mod)["main"].body
- tvm.ir.assert_structural_equal(new_stmt, stmt)
-
- with tvm.transform.PassContext(
- config={"s_tir.HoistIfThenElse": {"support_block_scope_hoisting":
True}}
- ):
- new_stmt = tvm.s_tir.transform.HoistIfThenElse()(mod)["main"].body
- tvm.ir.assert_structural_equal(new_stmt, stmt)
-
-
-def test_no_hoisting_2():
- @Ts.function(private=True)
- def func(data: T.handle("float32")):
- data_ptr = T.decl_tensor(1, "float32", data=data)
- for i in range(10):
- for j in range(10):
- for k in range(10):
- if i <= 3:
- data_ptr[i * 100 + j * 10 + k] = data_ptr[i * 100 + j
* 10 + k] + T.float32(
- 0.3
- )
- data_ptr[i * 100 + j * 10 + k] = data_ptr[i * 100 + j * 10
+ k] + T.float32(0.5)
-
- mod = tvm.IRModule.from_expr(func)
- stmt = mod["main"].body
- new_stmt = tvm.s_tir.transform.HoistIfThenElse()(mod)["main"].body
- tvm.ir.assert_structural_equal(new_stmt, stmt)
-
- with tvm.transform.PassContext(
- config={"s_tir.HoistIfThenElse": {"support_block_scope_hoisting":
True}}
- ):
- new_stmt = tvm.s_tir.transform.HoistIfThenElse()(mod)["main"].body
- tvm.ir.assert_structural_equal(new_stmt, stmt)
-
-
-def test_no_hoisting_4():
- dshape = (32, 64)
- dshape_inner = (33, 63)
-
- @I.ir_module
- class Module:
- @Ts.function(private=True)
- def main(data: T.Tensor((1,), "float32"), l: T.int32, m: T.int32, n:
T.int32):
- bx = T.launch_thread("blockIdx.x", dshape[1])
- for i in T.serial(l):
- for j in T.serial(m):
- for k in T.serial(n):
- tx_var = T.launch_thread("threadIdx.x",
dshape_inner[0])
- if tx_var < 3:
- data[bx * j + tx_var * j * k] = data[
- bx * j + tx_var * j * k
- ] + T.float32(0.3)
- else:
- data[bx * j + tx_var * j * k] = data[
- bx * j + tx_var * j * k
- ] + T.float32(1.3)
-
- stmt = Module["main"].body
- new_stmt = tvm.s_tir.transform.HoistIfThenElse()(Module)["main"].body
- tvm.ir.assert_structural_equal(new_stmt, stmt)
-
- with tvm.transform.PassContext(
- config={"s_tir.HoistIfThenElse": {"support_block_scope_hoisting":
True}}
- ):
- new_stmt = tvm.s_tir.transform.HoistIfThenElse()(Module)["main"].body
- tvm.ir.assert_structural_equal(new_stmt, stmt)
-
-
-def test_no_hoisting_6():
- dshape = (32, 64)
-
- @I.ir_module
- class Module:
- @Ts.function(private=True)
- def main(data: T.Tensor((1,), "float32"), l: T.int32, m: T.int32, n:
T.int32):
- tx = T.launch_thread("threadIdx.x", dshape[0])
- bx = T.launch_thread("blockIdx.x", dshape[1])
- for i in T.serial(l):
- for j in T.serial(m):
- for k in T.serial(n):
- if tx + k < 3:
- data[bx * j + tx * j * k] = data[bx * j + tx * j *
k] + T.float32(0.3)
- else:
- data[bx * j + tx * j * k] = data[bx * j + tx * j *
k] + T.float32(1.3)
-
- stmt = Module["main"].body
- new_stmt = tvm.s_tir.transform.HoistIfThenElse()(Module)["main"].body
- tvm.ir.assert_structural_equal(new_stmt, stmt)
-
- with tvm.transform.PassContext(
- config={"s_tir.HoistIfThenElse": {"support_block_scope_hoisting":
True}}
- ):
- new_stmt = tvm.s_tir.transform.HoistIfThenElse()(Module)["main"].body
- tvm.ir.assert_structural_equal(new_stmt, stmt)
-
-
-def test_no_hoisting_7():
- dshape = (32, 64)
-
- @I.ir_module
- class Module:
- @Ts.function(private=True)
- def main(data: T.Tensor((1,), "float32"), l: T.int32, m: T.int32, n:
T.int32):
- tx = T.launch_thread("threadIdx.x", dshape[0])
- bx = T.launch_thread("blockIdx.x", dshape[1])
- for i in T.serial(l):
- for j in T.serial(m):
- if tx + j < 9:
- for k in T.serial(n):
- if tx + k < 3:
- data[bx * j + tx * j * k] = data[bx * j + tx *
j * k] + T.float32(
- 0.3
- )
-
- stmt = Module["main"].body
- new_stmt = tvm.s_tir.transform.HoistIfThenElse()(Module)["main"].body
- tvm.ir.assert_structural_equal(new_stmt, stmt)
-
- with tvm.transform.PassContext(
- config={"s_tir.HoistIfThenElse": {"support_block_scope_hoisting":
True}}
- ):
- new_stmt = tvm.s_tir.transform.HoistIfThenElse()(Module)["main"].body
- tvm.ir.assert_structural_equal(new_stmt, stmt)
-
-
-def test_hoisting_block_scope_2():
- dshape = (32, 64)
-
- @I.ir_module
- class Module:
- @Ts.function(private=True)
- def main(data: T.Tensor((1,), "float32"), l: T.int32, m: T.int32, n:
T.int32):
- tx = T.launch_thread("threadIdx.x", dshape[0])
- for i in T.serial(l):
- for j in T.serial(m):
- for k in T.serial(n):
- bx_var = T.launch_thread("blockIdx.x", dshape[1])
- if tx < 3:
- data[bx_var * j + tx * j * k] = data[
- bx_var * j + tx * j * k
- ] + T.float32(0.3)
- else:
- data[bx_var * j + tx * j * k] = data[
- bx_var * j + tx * j * k
- ] + T.float32(1.3)
-
- mod = Module
- mod = tvm.tirx.transform.StmtSimplify()(mod)
- mod = tvm.tirx.transform.RemoveNoOp()(mod)
- stmt = mod["main"].body
-
- new_stmt = tvm.s_tir.transform.HoistIfThenElse()(mod)["main"].body
- tvm.ir.assert_structural_equal(new_stmt, stmt)
-
- with tvm.transform.PassContext(
- config={"s_tir.HoistIfThenElse": {"support_block_scope_hoisting":
True}}
- ):
- new_stmt = tvm.s_tir.transform.HoistIfThenElse()(mod)["main"].body
- assert not tvm_ffi.structural_equal(new_stmt, stmt)
-
-
-def test_hoisting_block_scope_5():
- @I.ir_module
- class Module:
- @Ts.function(private=True)
- def main(data: T.Tensor((1,), "float32"), l: T.int32, m: T.int32, n:
T.int32, g: T.int32):
- for i in T.serial(l):
- for j in T.serial(m):
- for k in T.serial(n):
- if data[g] < T.float32(3):
- data[9 * j + 3 * j * k] = data[9 * j + 3 * j * k]
+ T.float32(0.3)
- else:
- data[9 * j + 3 * j * k] = data[9 * j + 3 * j * k]
+ T.float32(1.3)
-
- stmt = Module["main"].body
- new_stmt = tvm.s_tir.transform.HoistIfThenElse()(Module)["main"].body
- assert not tvm_ffi.structural_equal(new_stmt, stmt)
-
- mod = tvm.IRModule.from_expr(tvm.tirx.Function([], new_stmt))
- stmt = new_stmt
-
- with tvm.transform.PassContext(
- config={"s_tir.HoistIfThenElse": {"support_block_scope_hoisting":
True}}
- ):
- new_stmt = tvm.s_tir.transform.HoistIfThenElse()(mod)["main"].body
- tvm.ir.assert_structural_equal(new_stmt, stmt)
-
-
-def test_hoisting_block_scope_6():
- dshape = (32, 64)
-
- @I.ir_module
- class Module:
- @Ts.function(private=True)
- def main(data: T.Tensor((1,), "float32"), l: T.int32, m: T.int32, n:
T.int32):
- tx = T.launch_thread("threadIdx.x", dshape[0])
- bx = T.launch_thread("blockIdx.x", dshape[1])
- for i in T.serial(l):
- for j in T.serial(m):
- for k in T.serial(n):
- if tx + n < 3:
- data[bx * j + tx * j * k] = data[bx * j + tx * j *
k] + T.float32(0.3)
- else:
- data[bx * j + tx * j * k] = data[bx * j + tx * j *
k] + T.float32(1.3)
-
- stmt = Module["main"].body
- new_stmt = tvm.s_tir.transform.HoistIfThenElse()(Module)["main"].body
- tvm.ir.assert_structural_equal(new_stmt, stmt)
-
- with tvm.transform.PassContext(
- config={"s_tir.HoistIfThenElse": {"support_block_scope_hoisting":
True}}
- ):
- new_stmt = tvm.s_tir.transform.HoistIfThenElse()(Module)["main"].body
- assert not tvm_ffi.structural_equal(new_stmt, stmt)
-
-
-def test_hoisting_block_scope_7():
- dshape = (32, 64)
-
- @I.ir_module
- class Module:
- @Ts.function(private=True)
- def main(data: T.Tensor((1,), "float32"), l: T.int32, m: T.int32, n:
T.int32):
- tx = T.launch_thread("threadIdx.x", dshape[0])
- bx = T.launch_thread("blockIdx.x", dshape[1])
- for i in T.serial(l):
- for j in T.serial(m):
- for k in T.serial(n):
- if tx + i < 3:
- data[bx * j + tx * j * k] = data[bx * j + tx * j *
k] + T.float32(0.3)
- else:
- data[bx * j + tx * j * k] = data[bx * j + tx * j *
k] + T.float32(1.3)
-
- stmt = Module["main"].body
- new_stmt = tvm.s_tir.transform.HoistIfThenElse()(Module)["main"].body
- tvm.ir.assert_structural_equal(new_stmt, stmt)
-
- with tvm.transform.PassContext(
- config={"s_tir.HoistIfThenElse": {"support_block_scope_hoisting":
True}}
- ):
- new_stmt = tvm.s_tir.transform.HoistIfThenElse()(Module)["main"].body
- assert not tvm_ffi.structural_equal(new_stmt, stmt)
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py
b/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py
index 5501667a35..82c0fd3d25 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py
@@ -20,6 +20,7 @@ import pytest
import tvm
import tvm.testing
from tvm import s_tir
+from tvm.script import ir as I
from tvm.script import s_tir as Ts
from tvm.script import tirx as T
from tvm.testing import env
@@ -237,6 +238,30 @@ def test_sync_bind():
tvm.ir.assert_structural_equal(mod["main"], expected)
[email protected]("loop", [T.serial, T.unroll])
+def test_hoist_if_preserves_sync_placement(loop):
+ @I.ir_module
+ class Before:
+ @T.function
+ def main(A: T.Tensor((4,), "int32")):
+ B = T.alloc_tensor((4,), "int32", scope="shared")
+ tx = T.launch_thread("threadIdx.x", 4)
+ B[tx] = A[tx]
+ for i in loop(1, 4):
+ if tx == 0:
+ B[i] = B[i - 1] + 1
+ A[tx] = B[tx]
+
+ # Each iteration must keep its own guard so synchronization can be inserted
+ # between shared-memory accesses, outside conditional control flow.
+ finish = tvm.transform.Sequential(
+ [tvm.tirx.transform.UnrollLoop(),
tvm.s_tir.transform.ThreadSync("shared")]
+ )
+ Expected = finish(Before)
+ After = finish(tvm.tirx.transform.HoistIf()(Before))
+ tvm.ir.assert_structural_equal(After, Expected)
+
+
if __name__ == "__main__":
test_thread_storage_sync()
test_sync_else_branch()
diff --git a/tests/python/tirx/transform/test_transform_hoist_if.py
b/tests/python/tirx/transform/test_transform_hoist_if.py
new file mode 100644
index 0000000000..cd0fb8cd5f
--- /dev/null
+++ b/tests/python/tirx/transform/test_transform_hoist_if.py
@@ -0,0 +1,423 @@
+# 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
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+
+def test_outermost_independent_loop():
+ @I.ir_module
+ class Before:
+ @T.function
+ def main(A: T.Tensor((4, 4, 4), "int32")):
+ for i in T.serial(4):
+ for j in T.serial(4):
+ for k in T.serial(4):
+ if i < 2:
+ A[i, j, k] = 1
+ else:
+ A[i, j, k] = 2
+
+ @I.ir_module
+ class Expected:
+ @T.function
+ def main(A: T.Tensor((4, 4, 4), "int32")):
+ for i in T.serial(4):
+ if i < 2:
+ for j in T.serial(4):
+ for k in T.serial(4):
+ A[i, j, k] = 1
+ else:
+ for j in T.serial(4):
+ for k in T.serial(4):
+ A[i, j, k] = 2
+
+ After = tvm.tirx.transform.HoistIf()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_multiple_loop_dependencies():
+ @I.ir_module
+ class Before:
+ @T.function
+ def main(A: T.Tensor((4, 4, 4), "int32")):
+ for i in T.serial(4):
+ for j in T.serial(4):
+ for k in T.serial(4):
+ if i + j < 2:
+ A[i, j, k] = 1
+
+ @I.ir_module
+ class Expected:
+ @T.function
+ def main(A: T.Tensor((4, 4, 4), "int32")):
+ for i in T.serial(4):
+ for j in T.serial(4):
+ if i + j < 2:
+ for k in T.serial(4):
+ A[i, j, k] = 1
+
+ After = tvm.tirx.transform.HoistIf()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_nested_conditions():
+ @I.ir_module
+ class Before:
+ @T.function
+ def main(A: T.Tensor((4, 4), "int32"), n: T.int32):
+ for i in T.serial(4):
+ if i < 2:
+ for j in T.serial(4):
+ if 0 < n:
+ A[i, j] = 1
+
+ @I.ir_module
+ class Expected:
+ @T.function
+ def main(A: T.Tensor((4, 4), "int32"), n: T.int32):
+ if 0 < n:
+ for i in T.serial(4):
+ if i < 2:
+ for j in T.serial(4):
+ A[i, j] = 1
+
+ # One-sided conditions must not clone the loop before cleanup: a chain of
+ # such conditions would otherwise produce exponentially many loop copies.
+ inserted = tvm.tirx.transform.HoistIf().passes[0](Before)
+ assert inserted["main"].body[0].else_case is None
+ After = tvm.tirx.transform.HoistIf()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_dependent_condition():
+ @I.ir_module
+ class Before:
+ @T.function
+ def main(A: T.Tensor((4,), "int32")):
+ for i in T.serial(4):
+ if i < 2:
+ A[i] = 1
+
+ Expected = Before
+ After = tvm.tirx.transform.HoistIf()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_sibling_statements():
+ @I.ir_module
+ class Before:
+ @T.function
+ def main(A: T.Tensor((4, 4), "int32"), n: T.int32):
+ for i in T.serial(4):
+ A[i, 0] = 0
+ for j in T.serial(4):
+ if 0 < n:
+ A[i, j] = 1
+
+ @I.ir_module
+ class Expected:
+ @T.function
+ def main(A: T.Tensor((4, 4), "int32"), n: T.int32):
+ for i in T.serial(4):
+ A[i, 0] = 0
+ if 0 < n:
+ for j in T.serial(4):
+ A[i, j] = 1
+
+ After = tvm.tirx.transform.HoistIf()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_binding_scope():
+ @I.ir_module
+ class Before:
+ @T.function
+ def main(A: T.Tensor((4, 4), "int32")):
+ for i in T.serial(4):
+ x: T.let[T.int32] = A[i, 0]
+ for j in T.serial(4):
+ if 0 < x:
+ A[i, j] = 1
+
+ @I.ir_module
+ class Expected:
+ @T.function
+ def main(A: T.Tensor((4, 4), "int32")):
+ for i in T.serial(4):
+ x: T.let[T.int32] = A[i, 0]
+ if 0 < x:
+ for j in T.serial(4):
+ A[i, j] = 1
+
+ After = tvm.tirx.transform.HoistIf()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_memory_condition():
+ @I.ir_module
+ class Before:
+ @T.function
+ def main(A: T.Tensor((4,), "int32")):
+ for i in T.serial(4):
+ if 0 < A[0]:
+ A[i] = 1
+
+ Expected = Before
+ After = tvm.tirx.transform.HoistIf()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_effectful_condition():
+ @I.ir_module
+ class Before:
+ @T.function
+ def main(A: T.Tensor((4,), "int32")):
+ for i in T.serial(4):
+ if 0 < T.call_extern("predicate", ty="int32"):
+ A[i] = 1
+
+ Expected = Before
+ After = tvm.tirx.transform.HoistIf()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_thread_region_boundary():
+ @I.ir_module
+ class Before:
+ @T.function
+ def main(A: T.Tensor((4, 4), "int32"), n: T.int32):
+ for i in T.serial(4):
+ with T.launch_thread("threadIdx.x", 4) as tx:
+ for j in T.serial(4):
+ if 0 < n:
+ A[i, tx] = j
+
+ Expected = Before
+ After = tvm.tirx.transform.HoistIf()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_guarded_division():
+ @I.ir_module
+ class Before:
+ @T.function
+ def main(A: T.Tensor((4,), "int32"), n: T.int32):
+ for i in T.serial(4):
+ if i < n:
+ if 1 < 100 // n:
+ A[i] = 1
+
+ Expected = Before
+ After = tvm.tirx.transform.HoistIf()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_zero_trip_division():
+ @I.ir_module
+ class Before:
+ @T.function
+ def main(A: T.Tensor((4,), "int32"), n: T.int32):
+ for i in T.serial(n):
+ if 1 < 100 // n:
+ A[i] = 1
+
+ Expected = Before
+ After = tvm.tirx.transform.HoistIf()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_enclosing_else_is_preserved():
+ @I.ir_module
+ class Before:
+ @T.function
+ def main(A: T.Tensor((4,), "int32"), n: T.int32):
+ for i in T.serial(4):
+ if i < 2:
+ if 0 < n:
+ A[i] = 1
+ else:
+ A[i] = 2
+
+ @I.ir_module
+ class Expected:
+ @T.function
+ def main(A: T.Tensor((4,), "int32"), n: T.int32):
+ if 0 < n:
+ for i in T.serial(4):
+ if i < 2:
+ A[i] = 1
+ else:
+ A[i] = 2
+ else:
+ for i in T.serial(4):
+ if not i < 2:
+ A[i] = 2
+
+ After = tvm.tirx.transform.HoistIf()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_effectful_enclosing_condition():
+ @I.ir_module
+ class Before:
+ @T.function
+ def main(A: T.Tensor((4,), "int32"), n: T.int32):
+ for i in T.serial(4):
+ if 0 < i + T.call_extern("predicate", ty="int32"):
+ if 0 < n:
+ A[i] = 1
+
+ Expected = Before
+ After = tvm.tirx.transform.HoistIf()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+
[email protected]("field", ["min", "extent", "step"])
+def test_effectful_loop_header(field):
+ value = tvm.tirx.call_extern("int32", "loop_header")
+ start = value if field == "min" else 0
+ stop = value if field == "extent" else 4
+ step = value if field == "step" else 1
+
+ @I.ir_module
+ class Before:
+ @T.function
+ def main(A: T.Tensor((4,), "int32"), n: T.int32):
+ for i in T.serial(start, stop, step=step):
+ if 0 < n:
+ A[i] = 1
+
+ Expected = Before
+ After = tvm.tirx.transform.HoistIf()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+
[email protected]("loop", [T.parallel, T.vectorized, T.unroll])
+def test_non_default_loop_boundary(loop):
+ @I.ir_module
+ class Before:
+ @T.function
+ def main(A: T.Tensor((4, 4, 4), "int32"), n: T.int32):
+ for i in T.serial(4):
+ for j in loop(4):
+ for k in T.serial(4):
+ if 0 < n:
+ A[i, j, k] = 1
+
+ @I.ir_module
+ class Expected:
+ @T.function
+ def main(A: T.Tensor((4, 4, 4), "int32"), n: T.int32):
+ for i in T.serial(4):
+ for j in loop(4):
+ if 0 < n:
+ for k in T.serial(4):
+ A[i, j, k] = 1
+
+ After = tvm.tirx.transform.HoistIf()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_nested_else_splits_once():
+ @I.ir_module
+ class Before:
+ @T.function
+ def main(A: T.Tensor((4,), "int32"), n: T.int32, m: T.int32):
+ for i in T.serial(4):
+ if 0 < n:
+ if 0 < m:
+ A[i] = 1
+ else:
+ A[i] = 2
+ else:
+ A[i] = 3
+
+ @I.ir_module
+ class Expected:
+ @T.function
+ def main(A: T.Tensor((4,), "int32"), n: T.int32, m: T.int32):
+ if 0 < n:
+ for i in T.serial(4):
+ if 0 < m:
+ A[i] = 1
+ else:
+ A[i] = 2
+ else:
+ for i in T.serial(4):
+ A[i] = 3
+
+ After = tvm.tirx.transform.HoistIf()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_nested_loop_splits_once():
+ @I.ir_module
+ class Before:
+ @T.function
+ def main(A: T.Tensor((4, 4, 4), "int32")):
+ for i in T.serial(4):
+ for j in T.serial(4):
+ for k in T.serial(4):
+ if i < 2:
+ A[i, j, k] = 1
+ elif j < 2:
+ A[i, j, k] = 2
+ else:
+ A[i, j, k] = 3
+
+ @I.ir_module
+ class Expected:
+ @T.function
+ def main(A: T.Tensor((4, 4, 4), "int32")):
+ for i in T.serial(4):
+ if i < 2:
+ for j in T.serial(4):
+ for k in T.serial(4):
+ A[i, j, k] = 1
+ else:
+ for j in T.serial(4):
+ for k in T.serial(4):
+ if j < 2:
+ A[i, j, k] = 2
+ else:
+ A[i, j, k] = 3
+
+ After = tvm.tirx.transform.HoistIf()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_thread_binding_boundary():
+ @I.ir_module
+ class Before:
+ @T.function
+ def main(A: T.Tensor((4, 4), "int32"), n: T.int32):
+ for tx in T.thread_binding(4, thread="threadIdx.x"):
+ for i in T.serial(4):
+ if 0 < n:
+ A[tx, i] = 1
+
+ Expected = Before
+ After = tvm.tirx.transform.HoistIf()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
