tvm icon indicating copy to clipboard operation
tvm copied to clipboard

Published 20 hours ago • verified publisher icon apache

[TVMScript] Support continue and break in tvmscript

Open wrongtest-intellif opened this issue 7 months ago • 6 comments

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 https://github.com/apache/tvm/pull/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.

wrongtest-intellif avatar Apr 03 '25 05:04 wrongtest-intellif

cc @Kathryn-cat @spectrometerHBH

tqchen avatar Apr 14 '25 11:04 tqchen

@tvm-bot re-run

wrongtest-intellif avatar Apr 16 '25 07:04 wrongtest-intellif

the issue seems to be pr title

tqchen avatar Apr 16 '25 21:04 tqchen

cc @tqchen @Kathryn-cat @spectrometerHBH

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

wrongtest-intellif avatar Apr 24 '25 03:04 wrongtest-intellif

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

Kathryn-cat avatar Apr 24 '25 15:04 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?

Kathryn-cat avatar Apr 24 '25 15:04 Kathryn-cat

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

FrozenGene avatar Sep 12 '25 04:09 FrozenGene

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.

wrongtest-intellif avatar Sep 16 '25 05:09 wrongtest-intellif

LGTM!

Kathryn-cat avatar Sep 18 '25 21:09 Kathryn-cat