This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tirx-standalone-root in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 4f60cde6c8b02687a04d2a20cea67a5c0030238f 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:
