[TVMScript] Support continue and break in tvmscript

[wrongtest-intellif]

Add TIR intrinsics for break and continue semantics. It aims to provide roundtrip support for TVMScript frontend and help users to develop complex algorithms in more pythonic fashion.

Below are some notes:

1. Like tir.ret introduced by [TIR] Support Return in TIR #7084, they are not first-class IR node but only intrinsic calls.

2. The change implement the codegen for c and llvm. In llvm target, we check the innermost serial loop scope must exists.

3. The change do not check the validity of TIR local function scopes and parallel launch scopes. Thus currently if the continue and break are misused in a parallel for scope, the error might only raised in codegen phase lately.

4. We have not implement precise TIR analysis of control flow with loop jumps, thus if there are ret, break, continue in loops，existing analysis (eg, region lowerbound & upperbound) might not be robust.

[tqchen]

cc @Kathryn-cat @spectrometerHBH

[wrongtest-intellif]

@tvm-bot re-run

[tqchen]

the issue seems to be pr title

[Copilot]

Pull Request Overview

This PR adds intrinsic support for break and continue statements in TVMScript to enable more natural Pythonic loop control in TIR.

• Implements TIR intrinsic functions and registers them in C++ and Python.
• Adds code generation support for C, LLVM, and CPU backends.
• Extends TVMScript parser and IR builder to handle break and continue.

Reviewed Changes

Copilot reviewed 16 out of 16 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/python/tvmscript/test_tvmscript_syntax_sugar.py	Adds tests for break and continue syntax conversions.
tests/python/tir-base/test_tir_base.py	Introduces tests for loop jump semantics with break and continue.
src/tir/op/op.cc	Adds intrinsic calls for break and continue.
src/tir/op/builtin.cc	Defines builtin functions for break and continue.
src/tir/ir/stmt.cc	Removes a check in While node to allow non-trivial conditions.
src/target/source/codegen_c.cc	Updates C codegen to translate intrinsic calls to 'break;' and 'continue;'.
src/target/llvm/codegen_llvm.h	Introduces loop frame management APIs for break and continue.
src/target/llvm/codegen_llvm.cc	Implements LLVM code generation for break and continue.
src/target/llvm/codegen_cpu.cc	Updates CPU backend to handle loop frame swapping for break and continue.
python/tvm/tir/op.py	Exposes Python APIs for the new intrinsics.
python/tvm/tir/init.py	Updates module init with new intrinsic function imports.
python/tvm/script/parser/tir/parser.py	Adds parser handlers for break and continue tokens.
python/tvm/script/parser/core/parser.py	Adds general visiting methods for break and continue nodes.
python/tvm/script/ir_builder/tir/ir.py	Exposes intrinsic functions in the TIR IR builder.
include/tvm/tir/op.h	Declares prototypes for break and continue intrinsics.
include/tvm/tir/builtin.h	Declares builtin operators for break and continue.

Comments suppressed due to low confidence (1)

python/tvm/script/parser/tir/parser.py:563

• The docstring for visit_break incorrectly states "continue"; it should be updated to "The break visiting method for tir."

    """The continue visiting method for tir."""

[wrongtest-intellif]

cc @tqchen @Kathryn-cat @spectrometerHBH

The pipeline get passed. Could you help to review? thanks!

[Kathryn-cat]

would you like to add tests to printer/parser as well?

[Kathryn-cat]

I think overall LGTM, the llvm codegen seems good, perhaps we can strengthen the test cases for "nested for/while loop" and "mixed break/continue in a loop" to make sure the control flow is correct?

[FrozenGene]

@wrongtest-intellif could we rebase the latest branch? i think this is also a very good feature for other dsl development based on tir

[spectrometerHBH]

Overall LGTM. One issue is that the lowering transformations might eliminate unit loops. With break/continue introduced, such loop elimination should be handled carefully

[wrongtest-intellif]

cc @spectrometerHBH @Kathryn-cat Very thanks for the suggestions!

"mixed break/continue in a loop"

• Here is a bundle of ai-gen nested cases to check primfunc behave the same with raw golden python.
  • https://gist.github.com/wrongtest-intellif/1c74cc7acbc48e34311cb6efc69eaaab

With break/continue introduced, such loop elimination should be handled carefully

• Add an annotation pass to help existing loop transformations be aware of loop jumps.

[Kathryn-cat]

LGTM!