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 218743c339 [REFACTOR][s_tir] Keep partition hints on their owning
loops (#20551)
218743c339 is described below
commit 218743c3390f7bbae273573d575a3de7039ddc29
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Oct 5 13:49:27 2026 -0400
[REFACTOR][s_tir] Keep partition hints on their owning loops (#20551)
Loop partition hints currently become enclosing attributes that point to
the variable defined by an inner loop. Keep this policy in a
`loop_partition_hint` For annotation and have LoopPartition consume the
owning loop directly, including expression-valued hints and recursive
partitions. Preserve unrelated annotations when splitting or unrolling
loops.
Run LoopPartition before opaque lowering in the existing partitioning
pipelines and handle thread-binding For nodes with the existing
hardware-thread partition rules. Custom pipelines with hinted
thread-binding loops must use the same order; opaque lowering reports an
ordering error instead of discarding the hint. In-tree producers use the
new annotation spelling.
---
include/tvm/s_tir/stmt.h | 4 +-
include/tvm/s_tir/transform.h | 1 +
include/tvm/tirx/transform.h | 1 +
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 | 6 +
python/tvm/tirx/transform/transform.py | 3 +
src/s_tir/transform/annotate_irregular_loop.cc | 2 +-
src/s_tir/transform/loop_partition.cc | 132 +++++++++++++++------
src/s_tir/transform/lower_opaque_block.cc | 4 +
src/tirx/transform/lower_tirx_opaque.cc | 4 +
.../test_s_tir_transform_compact_buffer_region.py | 4 +-
...est_s_tir_transform_inject_software_pipeline.py | 4 +-
.../test_s_tir_transform_loop_partition.py | 38 +++---
15 files changed, 143 insertions(+), 66 deletions(-)
diff --git a/include/tvm/s_tir/stmt.h b/include/tvm/s_tir/stmt.h
index 5e4978d24c..47fcbd0765 100644
--- a/include/tvm/s_tir/stmt.h
+++ b/include/tvm/s_tir/stmt.h
@@ -254,9 +254,9 @@ constexpr const char* fragment_shape = "fragment_shape";
constexpr const char* fragment_layout = "fragment_layout";
/*!
- * \brief Mark that the loop should be partitioned.
+ * \brief For annotation requesting partitioning when its PrimExpr value is
provably true.
*/
-constexpr const char* pragma_loop_partition_hint =
"pragma_loop_partition_hint";
+constexpr const char* loop_partition_hint = "loop_partition_hint";
/*! \brief Mark of reduce scope */
constexpr const char* reduce_scope = "reduce_scope";
diff --git a/include/tvm/s_tir/transform.h b/include/tvm/s_tir/transform.h
index 0f2c2c4b1f..90cdf5cba1 100644
--- a/include/tvm/s_tir/transform.h
+++ b/include/tvm/s_tir/transform.h
@@ -158,6 +158,7 @@ TVM_DLL Pass TransformMmaBufferLayout();
/*!
* \brief Remove the block to ensure that the TIR can not be scheduled again.
+ * Run LoopPartition first when thread-binding loops carry loop_partition_hint.
* \return The pass.
*/
TVM_DLL Pass LowerOpaqueBlock();
diff --git a/include/tvm/tirx/transform.h b/include/tvm/tirx/transform.h
index dbbfe60689..8e95ffb7fd 100644
--- a/include/tvm/tirx/transform.h
+++ b/include/tvm/tirx/transform.h
@@ -345,6 +345,7 @@ TVM_DLL Pass LowerTIRxCleanup();
* \brief Lower opaque constructs in TIRX programs: allocation calls,
For(thread_binding),
* unit loop elimination. This is the tirx-specific counterpart of
* s_tir::LowerOpaqueBlock, without any SBlock handling.
+ * Run LoopPartition first when thread-binding loops carry loop_partition_hint.
* \return The pass.
*/
TVM_DLL Pass LowerTIRxOpaque();
diff --git a/python/tvm/backend/trn/pipeline.py
b/python/tvm/backend/trn/pipeline.py
index d0663218e6..e0225400bf 100644
--- a/python/tvm/backend/trn/pipeline.py
+++ b/python/tvm/backend/trn/pipeline.py
@@ -36,8 +36,8 @@ def trn_pipeline():
trn_transform.LowerTrainiumLayout(),
tvm.s_tir.transform.DecorateDeviceScope(),
tirx.transform.StmtSimplify(),
- tirx.transform.LowerTIRxOpaque(),
tvm.s_tir.transform.LoopPartition(),
+ tirx.transform.LowerTIRxOpaque(),
tvm.s_tir.transform.HoistIfThenElse(),
tirx.transform.StmtSimplify(),
tirx.transform.RemoveNoOp(),
diff --git a/python/tvm/s_tir/backend/adreno/pipeline.py
b/python/tvm/s_tir/backend/adreno/pipeline.py
index 32f486d4de..e5a8dd0774 100644
--- a/python/tvm/s_tir/backend/adreno/pipeline.py
+++ b/python/tvm/s_tir/backend/adreno/pipeline.py
@@ -49,12 +49,12 @@ def default_tir_pipeline(*, prepare_only=False):
s_tir.transform.AnnotateIrregularLoop(),
s_tir.transform.InjectSoftwarePipeline(),
s_tir.transform.TransformMmaBufferLayout(),
+ s_tir.transform.LoopPartition(),
s_tir.transform.LowerOpaqueBlock(),
s_tir.backend.adreno.transform.InjectTextureAlloc(),
tirx.transform.FlattenBuffer(),
tirx.transform.BF16ComputeLegalize(),
tirx.transform.NarrowDataType(32),
- s_tir.transform.LoopPartition(),
tirx.transform.VectorizeLoop(not
bool(config.get("tirx.disable_vectorize", False))),
s_tir.transform.InjectVirtualThread(),
s_tir.transform.InjectDoubleBuffer(),
diff --git a/python/tvm/s_tir/pipeline.py b/python/tvm/s_tir/pipeline.py
index ee069b4fb8..d6f159f976 100644
--- a/python/tvm/s_tir/pipeline.py
+++ b/python/tvm/s_tir/pipeline.py
@@ -50,11 +50,11 @@ def default_s_tir_pipeline(*, prepare_only=False):
s_tir.transform.AnnotateIrregularLoop(),
s_tir.transform.InjectSoftwarePipeline(),
s_tir.transform.TransformMmaBufferLayout(),
+ s_tir.transform.LoopPartition(),
s_tir.transform.LowerOpaqueBlock(),
tirx.transform.FlattenBuffer(),
tirx.transform.BF16ComputeLegalize(),
tirx.transform.NarrowDataType(32),
- s_tir.transform.LoopPartition(),
tirx.transform.VectorizeLoop(not
bool(config.get("tirx.disable_vectorize", False))),
s_tir.transform.InjectVirtualThread(),
s_tir.transform.InjectDoubleBuffer(),
diff --git a/python/tvm/s_tir/transform/transform.py
b/python/tvm/s_tir/transform/transform.py
index e6095a656b..65eb20b4f9 100644
--- a/python/tvm/s_tir/transform/transform.py
+++ b/python/tvm/s_tir/transform/transform.py
@@ -136,6 +136,8 @@ def LowerMatchBuffer():
def LowerOpaqueBlock():
"""Remove the block to ensure that the TIR can not be scheduled again.
+ Run LoopPartition first when thread-binding loops carry
loop_partition_hint.
+
Returns
-------
fpass : tvm.transform.Pass
@@ -230,6 +232,10 @@ class LoopPartitionConfig(_ffi.Object):
def LoopPartition():
"""Partition loops in the stmt.
+ Consumes the owning For's ``loop_partition_hint`` annotation. A provably
+ true value enables partitioning even without likely tags. Run this pass
+ before opaque lowering when hinting thread-binding loops.
+
Returns
-------
fpass : tvm.transform.Pass
diff --git a/python/tvm/tirx/transform/transform.py
b/python/tvm/tirx/transform/transform.py
index 57d46d03dd..06184633d9 100644
--- a/python/tvm/tirx/transform/transform.py
+++ b/python/tvm/tirx/transform/transform.py
@@ -549,6 +549,9 @@ def LowerTIRxOpaque():
This is the tirx-specific counterpart of s_tir.LowerOpaqueBlock,
without any SBlock/SBlockRealize handling.
+ Run s_tir.transform.LoopPartition first when thread-binding loops carry
+ loop_partition_hint.
+
Returns
-------
fpass : tvm.transform.Pass
diff --git a/src/s_tir/transform/annotate_irregular_loop.cc
b/src/s_tir/transform/annotate_irregular_loop.cc
index 3519271665..b56dd8a3d7 100644
--- a/src/s_tir/transform/annotate_irregular_loop.cc
+++ b/src/s_tir/transform/annotate_irregular_loop.cc
@@ -59,7 +59,7 @@ class IrregularLoopAnnotator : public StmtExprMutator {
<< "Loop kind " << op->kind << " is invalid for irregular loop " <<
op->loop_var;
for (const char* key :
{tirx::attr::pragma_auto_unroll_max_step,
tirx::attr::pragma_unroll_explicit,
- s_tir::attr::pragma_loop_partition_hint,
s_tir::attr::software_pipeline_stage}) {
+ s_tir::attr::loop_partition_hint,
s_tir::attr::software_pipeline_stage}) {
TVM_FFI_ICHECK(!res->annotations.count(key))
<< "Annotation `" << key << "` is invalid for irregular loop " <<
op->loop_var;
}
diff --git a/src/s_tir/transform/loop_partition.cc
b/src/s_tir/transform/loop_partition.cc
index be3561eb20..537225f41d 100644
--- a/src/s_tir/transform/loop_partition.cc
+++ b/src/s_tir/transform/loop_partition.cc
@@ -65,8 +65,7 @@ struct LoopPartitionConfigNode : public ffi::Object {
"Don't unroll loops with extent 1", refl::DefaultValue(false))
.def_ro("unroll_loop_with_partition_hint_no_interval",
&LoopPartitionConfigNode::unroll_loop_with_partition_hint_no_interval,
- "Unroll loops with pragma_loop_partition_hint and no interval",
- refl::DefaultValue(false));
+ "Unroll loops with loop_partition_hint and no interval",
refl::DefaultValue(false));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.transform.LoopPartitionConfig",
LoopPartitionConfigNode,
ffi::Object);
@@ -108,6 +107,46 @@ using Partition = std::unordered_map<PartitionKey, IntSet,
PartitionKeyHash, Par
using ExpressionSet = std::unordered_set<PrimExpr, ffi::ObjectPtrHash,
ffi::ObjectPtrEqual>;
+// Virtual threads are lowered separately and are not hardware partition
scopes.
+static bool IsVirtualThread(const ForNode* op) {
+ if (op->kind != ForKind::kThreadBinding) return false;
+ const auto& tag = op->thread_binding.value()->thread_tag;
+ return tag == "vthread" || tag == "vthread.x" || tag == "vthread.y" || tag
== "vthread.z";
+}
+
+static bool HasPartitionHint(const ForNode* op) {
+ auto it = op->annotations.find(attr::loop_partition_hint);
+ return it != op->annotations.end() &&
is_one((*it).second.as_or_throw<PrimExpr>());
+}
+
+// Resolve expression-valued hints before partitioning clones or renames loops.
+// Like the original hint selection, this proof does not assume loop bounds.
+class NormalizePartitionHints : public StmtExprMutator {
+ public:
+ using StmtExprMutator::Mutate;
+ using StmtExprMutator::Mutate_;
+
+ UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final
{
+ For loop = StmtExprMutator::Mutate_(op, inplace_mode)
+ .ValueOrUnchanged(ffi::GetRef<Stmt>(op))
+ .as_or_throw<For>();
+ auto it = loop->annotations.find(attr::loop_partition_hint);
+ if (it != loop->annotations.end()) {
+ auto value = (*it).second;
+ if (value == nullptr) {
+ loop.CopyOnWrite()->annotations.erase(attr::loop_partition_hint);
+ } else {
+ bool enabled = analyzer_->CanProve(value.cast<PrimExpr>());
+ loop.CopyOnWrite()->annotations.Set(attr::loop_partition_hint,
IntImm::Bool(enabled));
+ }
+ }
+ return loop;
+ }
+
+ private:
+ sym::Analyzer analyzer_;
+};
+
// Select potential candidate IRs that can be partitioned.
// Rule:
// - the range should not be const
@@ -120,12 +159,16 @@ class CandidateSelector final : public StmtExprVisitor {
: partition_const_loop_(partition_const_loop) {}
ffi::Optional<VisitInterrupt> Visit_(const ForNode* op) final {
- // always treat var with hint to be partitioned
+ if (IsVirtualThread(op)) return StmtExprVisitor::Visit_(op);
const VarNode* var = op->loop_var.get();
- if (partition_hint_vars.count(var)) {
+ if (HasPartitionHint(op)) {
candidates.insert(ffi::GetRef<Stmt>(op));
return StmtExprVisitor::Visit_(op);
}
+ if (op->kind == ForKind::kThreadBinding &&
+
runtime::ThreadScope::Create(op->thread_binding.value()->thread_tag).rank != 0)
{
+ return StmtExprVisitor::Visit_(op);
+ }
// partition const loop when sets partition_const_loop_
if (!is_const_int(op->min) || !is_const_int(op->extent) ||
partition_const_loop_) {
record_.insert({var, false});
@@ -145,11 +188,6 @@ class CandidateSelector final : public StmtExprVisitor {
const IterVarNode* iv = op->node.as<IterVarNode>();
TVM_FFI_ICHECK(iv);
Var var = iv->var;
- // always treat var with hint to be partitioned
- if (partition_hint_vars.count(var.get())) {
- candidates.insert(ffi::GetRef<Stmt>(op));
- return StmtExprVisitor::Visit_(op);
- }
runtime::ThreadScope scope =
runtime::ThreadScope::Create(iv->thread_tag);
auto value = op->value.as<PrimExpr>();
if ((scope.rank == 0) && (!value || !is_const_int(value.value()) ||
partition_const_loop_)) {
@@ -161,17 +199,6 @@ class CandidateSelector final : public StmtExprVisitor {
record_.erase(var.get());
return std::nullopt;
}
- } else if (op->attr_key == s_tir::attr::pragma_loop_partition_hint) {
- if (analyzer_->CanProve(op->value.as_or_throw<PrimExpr>())) {
- const VarNode* var = nullptr;
- if (op->node.as<VarNode>()) {
- var = op->node.as<VarNode>();
- } else if (op->node.as<IterVarNode>()) {
- var = op->node.as<IterVarNode>()->var.get();
- }
- TVM_FFI_ICHECK(var);
- partition_hint_vars.insert(var);
- }
}
return StmtExprVisitor::Visit_(op);
}
@@ -212,14 +239,12 @@ class CandidateSelector final : public StmtExprVisitor {
}
std::unordered_set<Stmt, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> candidates;
- std::unordered_set<const VarNode*> partition_hint_vars;
private:
bool in_likely_{false};
bool no_split_{false};
bool partition_const_loop_{false};
std::unordered_map<const VarNode*, VarIsUsed> record_;
- sym::Analyzer analyzer_;
};
// Finder try best to find partitions for hinted vars
@@ -254,6 +279,7 @@ class PartitionFinder : public StmtExprVisitor {
}
ffi::Optional<VisitInterrupt> Visit_(const ForNode* op) final {
+ if (IsVirtualThread(op)) return StmtExprVisitor::Visit_(op);
auto f_vset_contains = [this](const VarNode* var) { return
out_vars_.count(var); };
auto walkfn = [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
return f_vset_contains(var.get()) ?
ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
@@ -425,6 +451,25 @@ class ThreadPartitionInserter : public StmtExprMutator {
explicit ThreadPartitionInserter(const ExpressionSet& ps, PrimExpr cond)
: ps_(ps), cond_(cond), innermost_thread_scope_(false) {}
+ UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final
{
+ if (op->kind != ForKind::kThreadBinding || IsVirtualThread(op)) {
+ return StmtExprMutator::Mutate_(op, inplace_mode);
+ }
+ innermost_thread_scope_ = true;
+ For loop = StmtExprMutator::Mutate_(op, inplace_mode)
+ .ValueOrUnchanged(ffi::GetRef<Stmt>(op))
+ .as_or_throw<For>();
+ if (innermost_thread_scope_) {
+ Stmt body = loop->body;
+ Stmt simplified_body = ffi::make_object<ConditionEliminator>(ps_)
+ ->Mutate(body, InplaceMode::kDisallow)
+ .ValueOrUnchanged(body);
+ loop.CopyOnWrite()->body = IfThenElse(cond_, simplified_body, body);
+ }
+ innermost_thread_scope_ = false;
+ return loop;
+ }
+
UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode)
final {
if (op->attr_key == tirx::attr::thread_extent) {
innermost_thread_scope_ = true;
@@ -474,19 +519,24 @@ class LoopPartitioner : public StmtExprMutator {
}
UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final
{
+ if (IsVirtualThread(op)) return StmtExprMutator::Mutate_(op, inplace_mode);
+ bool thread_scope = op->kind == ForKind::kThreadBinding;
analyzer_->Bind(op->loop_var, Range::FromMinExtent(op->min, op->extent),
true);
auto fs = ffi::GetRef<Stmt>(op);
if (selector->candidates.count(fs)) {
- auto s = TryPartition(fs, op->loop_var, op->min, op->min + op->extent -
1, op->body, false);
+ auto s =
+ TryPartition(fs, op->loop_var, op->min, op->min + op->extent - 1,
op->body, thread_scope);
if (s.has_value()) return s.value();
}
- // normal path when loop partition fails
- // normal loop variable can be put into hint map.
- hint_map_.insert({op->loop_var.get(), IntSet::Interval(op->min, op->min +
op->extent - 1)});
+ // Relax threadIdx ranges to avoid introducing divergent partition
branches.
+ bool relax = thread_scope &&
+
runtime::ThreadScope::Create(op->thread_binding.value()->thread_tag).rank == 1;
+ auto& ranges = relax ? relax_map_ : hint_map_;
+ ranges.insert({op->loop_var.get(), IntSet::Interval(op->min, op->min +
op->extent - 1)});
Stmt res = StmtExprMutator::Mutate_(op, InplaceMode::kDisallow)
.ValueOrUnchanged(ffi::GetRef<Stmt>(op));
- hint_map_.erase(op->loop_var.get());
+ ranges.erase(op->loop_var.get());
return res;
}
@@ -647,7 +697,8 @@ ffi::Optional<Stmt> LoopPartitioner::TryPartition(const
Stmt& stmt, Var var, Pri
// include hint of var.
hint_map_.insert({var.get(), IntSet::Interval(min, max)});
- bool has_partition_hint_ = selector->partition_hint_vars.count(var.get());
+ const auto* loop = stmt.as<ForNode>();
+ bool has_partition_hint_ = loop && HasPartitionHint(loop);
auto finder = ffi::make_object<PartitionFinder>(var, hint_map_, relax_map_,
has_partition_hint_);
finder->Visit(body);
@@ -719,10 +770,13 @@ ffi::Optional<Stmt> LoopPartitioner::TryPartition(const
Stmt& stmt, Var var, Pri
}
if (!opt_cond_value.has_value()) {
- if (has_partition_hint_ && unroll_loop_with_partition_hint_no_interval_ &&
- analyzer_->CanProve(max - min > 0)) {
+ if (!partition_thread_scope && has_partition_hint_ &&
+ unroll_loop_with_partition_hint_no_interval_ &&
analyzer_->CanProve(max - min > 0)) {
auto new_body = VisitAndMutate(body);
- return For(var.as_or_throw<PrimVar>(), min, max - min + 1,
ForKind::kUnrolled, new_body);
+ auto unrolled = ffi::make_object<ForNode>(*loop);
+ unrolled->kind = ForKind::kUnrolled;
+ unrolled->body = std::move(new_body);
+ return For(unrolled);
}
return std::nullopt;
}
@@ -856,7 +910,7 @@ inline Stmt LoopPartitioner::MakeFor(const ffi::Object*
node, PrimExpr extent, S
TVM_FFI_ICHECK(for_node);
if (analyzer_->CanProve(extent == IntImm::Int32(1)) &&
!no_unroll_loop_with_extent_one_ &&
- for_node->annotations.empty()) {
+ for_node->annotations.size() ==
for_node->annotations.count(attr::loop_partition_hint)) {
// If the loop extent is 1, do not create the loop anymore
auto f_substitute = [loop_var = for_node->loop_var](
const Var& var) ->
ffi::Expected<ffi::UnchangedOr<ffi::Any>> {
@@ -891,16 +945,22 @@ class RemoveLikelyTagsAndHints : public StmtExprMutator {
}
}
- UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode)
final {
- if (op->attr_key == s_tir::attr::pragma_loop_partition_hint) {
- return Mutate(op->body, inplace_mode).ValueOrUnchanged(op->body);
+ UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final
{
+ For loop = StmtExprMutator::Mutate_(op, inplace_mode)
+ .ValueOrUnchanged(ffi::GetRef<Stmt>(op))
+ .as_or_throw<For>();
+ if (loop->annotations.count(attr::loop_partition_hint)) {
+ loop.CopyOnWrite()->annotations.erase(attr::loop_partition_hint);
}
- return StmtExprMutator::Mutate_(op, inplace_mode);
+ return loop;
}
};
Stmt LoopPartition(Stmt stmt, bool partition_const_loop, bool
no_unroll_loop_with_extent_one,
bool unroll_loop_with_partition_hint_no_interval) {
+ stmt = ffi::make_object<NormalizePartitionHints>()
+ ->Mutate(stmt, InplaceMode::kAllow)
+ .ValueOrUnchanged(std::move(stmt));
stmt = ffi::make_object<LoopPartitioner>(partition_const_loop,
no_unroll_loop_with_extent_one,
unroll_loop_with_partition_hint_no_interval)
->VisitAndMutate(std::move(stmt));
diff --git a/src/s_tir/transform/lower_opaque_block.cc
b/src/s_tir/transform/lower_opaque_block.cc
index bbb68f392e..4bf291ee38 100644
--- a/src/s_tir/transform/lower_opaque_block.cc
+++ b/src/s_tir/transform/lower_opaque_block.cc
@@ -115,6 +115,10 @@ class OpaqueBlockLower : public StmtExprMutator {
// Step 4. Create new For loop accordingly
if (op->kind == ForKind::kThreadBinding) {
// Case 1. Thread binding
+ TVM_FFI_ICHECK(!op->annotations.count("loop_partition_hint") ||
+ op->annotations.at("loop_partition_hint") == nullptr)
+ << "Run LoopPartition before opaque lowering of a thread-binding
loop with "
+ "loop_partition_hint";
TVM_FFI_ICHECK(op->thread_binding.has_value());
ffi::String thread_tag = op->thread_binding.value()->thread_tag;
body = MakeLaunchThread(min, extent, op->loop_var, thread_tag, body);
diff --git a/src/tirx/transform/lower_tirx_opaque.cc
b/src/tirx/transform/lower_tirx_opaque.cc
index 949d9d821d..d8792dda9d 100644
--- a/src/tirx/transform/lower_tirx_opaque.cc
+++ b/src/tirx/transform/lower_tirx_opaque.cc
@@ -72,6 +72,10 @@ class TIRxOpaqueLower : public StmtExprMutator {
// Step 4. Create new For loop accordingly
if (op->kind == ForKind::kThreadBinding) {
// Case 1. Thread binding → AttrStmt(thread_extent)
+ TVM_FFI_ICHECK(!op->annotations.count("loop_partition_hint") ||
+ op->annotations.at("loop_partition_hint") == nullptr)
+ << "Run LoopPartition before opaque lowering of a thread-binding
loop with "
+ "loop_partition_hint";
TVM_FFI_ICHECK(op->thread_binding.has_value());
ffi::String thread_tag = op->thread_binding.value()->thread_tag;
body = MakeLaunchThread(min, extent, op->loop_var, thread_tag, body);
diff --git
a/tests/python/s_tir/transform/test_s_tir_transform_compact_buffer_region.py
b/tests/python/s_tir/transform/test_s_tir_transform_compact_buffer_region.py
index a28c75387c..1760d81d04 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_compact_buffer_region.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_compact_buffer_region.py
@@ -1011,8 +1011,8 @@ class TestTileAwareCompaction(BaseCompactTest):
B: T.Tensor((128, 128), "float32"),
C: T.Tensor((128, 128), "float32"),
):
- for i_0 in range(5, annotations={"pragma_loop_partition_hint": 1}):
- for j_0 in range(5, annotations={"pragma_loop_partition_hint":
1}):
+ for i_0 in range(5, annotations={"loop_partition_hint": 1}):
+ for j_0 in range(5, annotations={"loop_partition_hint": 1}):
A_local = T.decl_tensor((26, 128), scope="local")
B_local = T.decl_tensor((128, 26), scope="local")
C_local = T.decl_tensor((26, 26), scope="local")
diff --git
a/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py
b/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py
index 438dcaa923..b13248c4e5 100644
---
a/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py
+++
b/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py
@@ -236,7 +236,7 @@ def simple_compute_with_other_annotation(
annotations={
"software_pipeline_stage": [0, 1],
"software_pipeline_order": [0, 1],
- "pragma_loop_partition_hint": True,
+ "loop_partition_hint": True,
},
):
with Ts.sblock("compute"):
@@ -272,7 +272,7 @@ def transformed_simple_compute_with_other_annotation(
for i in T.serial(
0,
15,
- annotations={"pragma_loop_partition_hint": True},
+ annotations={"loop_partition_hint": True},
):
with Ts.sblock():
Ts.reads([A[tx, i + 1]])
diff --git
a/tests/python/s_tir/transform/test_s_tir_transform_loop_partition.py
b/tests/python/s_tir/transform/test_s_tir_transform_loop_partition.py
index 25d48f99a8..2df8c3e69a 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_loop_partition.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_loop_partition.py
@@ -223,10 +223,10 @@ def partitioned_concat(
def partition_from_scheduled_tir(prim_func, pass_cfg, do_flatten=True):
with tvm.transform.PassContext(config=pass_cfg):
mod = IRModule.from_expr(prim_func.with_attr("global_symbol", "main"))
+ mod = tvm.s_tir.transform.LoopPartition()(mod)
mod = tvm.s_tir.transform.LowerOpaqueBlock()(mod)
if do_flatten:
mod = tvm.tirx.transform.FlattenBuffer()(mod)
- mod = tvm.s_tir.transform.LoopPartition()(mod)
mod = tvm.tirx.transform.StmtSimplify()(mod)
mod = tvm.tirx.transform.RemoveNoOp()(mod)
return mod
@@ -262,7 +262,7 @@ def concat_func_3(
placeholder_1_flat = T.decl_tensor([25088], "int8",
data=placeholder_1.data)
placeholder_2_flat = T.decl_tensor([25088], "int8",
data=placeholder_2.data)
T_concat_flat = T.decl_tensor([100352], "int8", data=T_concat.data)
- for i1 in T.serial(128, annotations={"pragma_loop_partition_hint": 1}):
+ for i1 in T.serial(128, annotations={"loop_partition_hint": 1}):
for i2, i3 in T.grid(28, 28):
if 96 <= i1:
T_concat_flat[i1 * 784 + i2 * 28 + i3] = placeholder_2_flat[
@@ -294,7 +294,7 @@ def test_loop_partition_unroll_hint():
B = T.decl_tensor(25088, "int8", data=B_arg.data)
for ax0 in T.serial(
112,
- annotations={"pragma_loop_partition_hint": True},
+ annotations={"loop_partition_hint": True},
):
for ax1, ax2, ax3 in T.grid(224, 7, 16):
if 3 <= ax0 * 2 + ax2 and ax0 * 2 + ax2 < 227 and ax3 < 3:
@@ -339,8 +339,8 @@ def test_loop_partition_recursive_unroll_hint():
@Ts.prim_func
def main():
placeholder_0_dm = T.decl_tensor([1, 32, 32, 16], dtype="int8")
- for i3_0 in T.serial(5, annotations={"pragma_loop_partition_hint": 1}):
- for i2_0 in T.serial(2, annotations={"pragma_loop_partition_hint":
1}):
+ for i3_0 in T.serial(5, annotations={"loop_partition_hint": 1}):
+ for i2_0 in T.serial(2, annotations={"loop_partition_hint": 1}):
pad_temp = T.decl_tensor([1, 16, 16, 16], dtype="int8")
for ax0, ax1, ax2 in T.grid(16, 16, 16):
if (
@@ -405,7 +405,7 @@ def test_loop_partition_keep_loop_annotations():
def before(A: T.Tensor(160, "int32"), B: T.Tensor(160, "int32")) -> None:
for i in T.serial(
160,
- annotations={"pragma_loop_partition_hint": True, "key": "value"},
+ annotations={"loop_partition_hint": True, "key": "value"},
):
if i < 10:
B[i] = A[i] + 1
@@ -445,7 +445,7 @@ def test_loop_partition_with_unit_loop_in_condition():
T_concat: T.Tensor((100352,), "int8"),
) -> None:
for k in range(1, annotations={"preserve_unit_loop": True}):
- for i1 in T.serial(128, annotations={"pragma_loop_partition_hint":
1}):
+ for i1 in T.serial(128, annotations={"loop_partition_hint": 1}):
for i2, i3 in T.grid(28, 28):
if 96 <= k * 128 + i1:
T_concat[k * i1 * 784 + i2 * 28 + i3] = placeholder_2[
@@ -498,7 +498,7 @@ def concat_func_single_point(
T_concat: T.Tensor((28, 128), "int8"),
) -> None:
for i0 in range(28):
- for i1 in T.serial(128, annotations={"pragma_loop_partition_hint": 1}):
+ for i1 in T.serial(128, annotations={"loop_partition_hint": 1}):
if i1 > 63:
T_concat[i0, i1] = placeholder[i0, i1 - 64]
elif i1 == 63:
@@ -534,7 +534,7 @@ def concat_func_start_point_equality(
T_concat: T.Tensor((28, 128), "int8"),
) -> None:
for i0 in range(28):
- for i1 in range(128, annotations={"pragma_loop_partition_hint": 1}):
+ for i1 in range(128, annotations={"loop_partition_hint": 1}):
if i1 == 0:
# Special case for i1 == 0
T_concat[i0, i1] = placeholder_1[i0, 0]
@@ -573,7 +573,7 @@ def concat_func_end_point_equality(
T_concat: T.Tensor((28, 128), "int8"),
) -> None:
for i0 in range(28):
- for i1 in range(128, annotations={"pragma_loop_partition_hint": 1}):
+ for i1 in range(128, annotations={"loop_partition_hint": 1}):
if i1 == 127:
# Explicit equality check for the end point i1 == 127
T_concat[i0, i1] = placeholder_1[i0, 0]
@@ -612,9 +612,7 @@ def concat_func_edge_equalities(
T_concat: T.Tensor((28, 66), "int8"),
) -> None:
for i0 in range(28):
- for i1 in range(
- 66, annotations={"pragma_loop_partition_hint": 1}
- ): # Loop from 0 to 65 inclusive
+ for i1 in range(66, annotations={"loop_partition_hint": 1}): # Loop
from 0 to 65 inclusive
if i1 == 0:
# Handle equality at the start of the range: i1 == 0
T_concat[i0, i1] = placeholder_2[i0, 0]
@@ -654,7 +652,7 @@ def concat_five_buffers_with_equalities(
T_concat: T.Tensor((28, 129), "int8"),
) -> None:
for i0 in range(28):
- for i1 in range(130, annotations={"pragma_loop_partition_hint": 1}):
+ for i1 in range(130, annotations={"loop_partition_hint": 1}):
if i1 == 0:
T_concat[i0, i1] = buffer_a[i0, 0]
elif i1 == 64:
@@ -694,13 +692,13 @@ def concat_five_buffers_with_equalities_expected(
@Ts.prim_func
def nested_partition_with_single_points(A: T.Tensor((25,), "int32")):
- for i in T.serial(5, annotations={"pragma_loop_partition_hint": 1}):
+ for i in T.serial(5, annotations={"loop_partition_hint": 1}):
if i == 1:
- for j in T.serial(5, annotations={"pragma_loop_partition_hint":
1}):
+ for j in T.serial(5, annotations={"loop_partition_hint": 1}):
if j > 2:
A[i * 5 + j] = i * 5 + j
else:
- for j in T.serial(5, annotations={"pragma_loop_partition_hint":
1}):
+ for j in T.serial(5, annotations={"loop_partition_hint": 1}):
if j > 2:
A[i * 5 + j] = i * 15 + j
@@ -745,9 +743,9 @@ def test_single_point_partition(origin, expected):
def test_equation_on_floordiv():
@Ts.prim_func
def before(A: T.Tensor((2, 2, 20), "int32")):
- for i in T.serial(5, annotations={"pragma_loop_partition_hint": 1}):
+ for i in T.serial(5, annotations={"loop_partition_hint": 1}):
if i == 1:
- for vv in T.vectorized(640,
annotations={"pragma_loop_partition_hint": 1}):
+ for vv in T.vectorized(640,
annotations={"loop_partition_hint": 1}):
if i * 2 + vv // 320 == 3:
A[i - 1, i * 2 + vv // 320 - 3, vv % 320 // 16] = 1
@@ -770,7 +768,7 @@ def test_ignore_loop_partition_hint():
def before(A: T.Tensor((10), "float32"), D: T.Tensor((10), "float32")):
B = T.decl_tensor([2], "float32")
C = T.decl_tensor([2], "float32")
- for i in T.serial(12, annotations={"pragma_loop_partition_hint": 1}):
+ for i in T.serial(12, annotations={"loop_partition_hint": 1}):
if T.ignore_loop_partition(i < 10):
B[i % 2] = A[i] + 1.0
if T.ignore_loop_partition(1 <= i and i < 11):
