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 069afb22bc [FIX][TIRx] Preserve loops containing break or continue
during unrolling (#20783)
069afb22bc is described below
commit 069afb22bc2e3d62f7d2c99295289beec30a091d
Author: Tianqi Chen <[email protected]>
AuthorDate: Sat Oct 17 17:44:41 2026 -0400
[FIX][TIRx] Preserve loops containing break or continue during unrolling
(#20783)
Preserve loops containing `break` or `continue` during TIRx unrolling so
loop-control statements retain their enclosing loop. Clear explicit
unroll markers on these loops and apply the guard before automatic,
local-access, and extent-one unrolling.
---
tirx/src/tirx/transform/unroll_loop.cc | 27 ++++++++++++++++
tirx/tests/python/tirx/transform/test_unroll.py | 41 +++++++++++++++++++++++++
2 files changed, 68 insertions(+)
diff --git a/tirx/src/tirx/transform/unroll_loop.cc
b/tirx/src/tirx/transform/unroll_loop.cc
index f1d7d89456..f5c8c2ac4c 100644
--- a/tirx/src/tirx/transform/unroll_loop.cc
+++ b/tirx/src/tirx/transform/unroll_loop.cc
@@ -130,6 +130,7 @@ class LoopUnroller : public StmtExprMutator {
UnchangedOr<Stmt> RewriteLoop(const ForNode* op, InplaceMode inplace_mode) {
// Post order so we can collect more information
+ int loop_control_count = loop_control_count_;
auto result = StmtExprMutator::Mutate_(op, inplace_mode);
if (!result.IsUnchanged()) {
op = ffi::AnyView(result).as<ForNode>();
@@ -151,6 +152,20 @@ class LoopUnroller : public StmtExprMutator {
node->annotations.erase(tvm::tirx::attr::kUnrollExplicit.name);
op = node;
}
+ // Keep the loop enclosing break/continue intact, including extent-one
loops.
+ if (loop_control_count_ != loop_control_count) {
+ ++normal_loop_depth_;
+ if (op->kind == ForKind::kUnrolled) {
+ if (inplace_mode == InplaceMode::kAllow) {
+ const_cast<ForNode*>(op)->kind = ForKind::kDefault;
+ return result;
+ }
+ auto node = ffi::make_object<ForNode>(*op);
+ node->kind = ForKind::kDefault;
+ return For(std::move(node));
+ }
+ return result;
+ }
int value = GetExtent(op);
// condition for auto unroll
bool auto_unroll = (op->kind == ForKind::kDefault && value >= 0 &&
normal_loop_depth_ == 0 &&
@@ -231,6 +246,16 @@ class LoopUnroller : public StmtExprMutator {
return StmtExprMutator::Mutate_(op, inplace_mode);
}
+ UnchangedOr<Stmt> Mutate_(const BreakNode* op, InplaceMode inplace_mode)
final {
+ ++loop_control_count_;
+ return ffi::Unchanged();
+ }
+
+ UnchangedOr<Stmt> Mutate_(const ContinueNode* op, InplaceMode inplace_mode)
final {
+ ++loop_control_count_;
+ return ffi::Unchanged();
+ }
+
UnchangedOr<Stmt> Mutate_(const EvaluateNode* op, InplaceMode inplace_mode)
final {
++step_count_;
return StmtExprMutator::Mutate_(op, inplace_mode);
@@ -311,6 +336,8 @@ class LoopUnroller : public StmtExprMutator {
int normal_loop_depth_{0};
// number of unrolled cases in current scope.
int unroll_depth_{0};
+ // Number of break/continue statements visited, including nested loops.
+ int loop_control_count_{0};
// Number of total steps unrolled
int step_count_{0};
// set of indices touched during visit local memory
diff --git a/tirx/tests/python/tirx/transform/test_unroll.py
b/tirx/tests/python/tirx/transform/test_unroll.py
index 32ddda054c..127ba27d49 100644
--- a/tirx/tests/python/tirx/transform/test_unroll.py
+++ b/tirx/tests/python/tirx/transform/test_unroll.py
@@ -196,6 +196,47 @@ def test_unroll_local_access():
tvm_tirx.ir.assert_structural_equal(after, Expected)
+def test_unroll_break_continue():
+ @I.ir_module
+ class Before:
+ @T.function
+ def main(A: T.Tensor((4,), "int32")):
+ for i in T.unroll(4):
+ if A[i] == 0:
+ break
+ A[i] = 1
+ for i in T.serial(1):
+ if A[i] == 0:
+ continue
+ A[i] = 2
+
+ @I.ir_module
+ class Expected:
+ @T.function
+ def main(A: T.Tensor((4,), "int32")):
+ for i in T.serial(4):
+ if A[i] == 0:
+ break
+ A[i] = 1
+ for i in T.serial(1):
+ if A[i] == 0:
+ continue
+ A[i] = 2
+
+ for explicit in (True, False):
+ with tvm_tirx.ir.transform.PassContext(
+ config={
+ "tirx.UnrollLoop": {
+ "auto_max_step": 16,
+ "auto_max_extent": 1,
+ "explicit_unroll": explicit,
+ }
+ }
+ ):
+ after = tvm_tirx.tirx.transform.UnrollLoop()(Before)
+ tvm_tirx.ir.assert_structural_equal(after, Expected)
+
+
def test_pass_config():
@I.ir_module
class Before:
