Skip to content

[TVMScript] Support continue and break in tvmscript #17804

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 4 commits into
base: main
Choose a base branch
from

Conversation

wrongtest-intellif
Copy link
Contributor

@wrongtest-intellif wrongtest-intellif commented Apr 3, 2025

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
Copy link
Member

tqchen commented Apr 14, 2025

cc @Kathryn-cat @spectrometerHBH

@wrongtest-intellif wrongtest-intellif force-pushed the support_continue_break_in_tvmscript branch from fde0281 to d0834fa Compare April 16, 2025 06:11
@wrongtest-intellif wrongtest-intellif marked this pull request as ready for review April 16, 2025 06:34
@wrongtest-intellif
Copy link
Contributor Author

wrongtest-intellif commented Apr 16, 2025

@tvm-bot re-run

@tqchen
Copy link
Member

tqchen commented Apr 16, 2025

the issue seems to be pr title

@wrongtest-intellif wrongtest-intellif force-pushed the support_continue_break_in_tvmscript branch from d0834fa to 5759756 Compare April 18, 2025 02:13
Copy link

@Copilot Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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."""

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
@wrongtest-intellif
Copy link
Contributor Author

cc @tqchen @Kathryn-cat @spectrometerHBH

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

@Kathryn-cat
Copy link
Contributor

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

@Kathryn-cat
Copy link
Contributor

Kathryn-cat commented Apr 24, 2025

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?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants