-
Notifications
You must be signed in to change notification settings - Fork 3.6k
[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
base: main
Are you sure you want to change the base?
[TVMScript] Support continue and break in tvmscript #17804
Conversation
fde0281
to
d0834fa
Compare
@tvm-bot re-run |
the issue seems to be pr title |
d0834fa
to
5759756
Compare
There was a problem hiding this 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>
cc @tqchen @Kathryn-cat @spectrometerHBH The pipeline get passed. Could you help to review? thanks! |
would you like to add tests to printer/parser as well? |
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? |
Add TIR intrinsics for
break
andcontinue
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:
Like
tir.ret
introduced by [TIR] Support Return in TIR #7084, they are not first-class IR node but only intrinsic calls.The change implement the codegen for
c
andllvm
. In llvm target, we check the innermost serial loop scope must exists.The change do not check the validity of TIR local function scopes and parallel launch scopes. Thus currently if the
continue
andbreak
are misused in a parallel for scope, the error might only raised in codegen phase lately.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.