[TVMScript] Support continue and break in tvmscript (#17804)

* support continue and break in tvmscript

* fix black format

* fix pylint issue

* Update tests/python/tvmscript/test_tvmscript_syntax_sugar.py

Co-authored-by: Copilot <[email protected]>

* add printer/parser test, fix lint

* Fit to latest ffi update

* Skip i386 numpy-related test

* Introduce AnnotateIrregularLoop before any lowering loop expansions.

---------

Co-authored-by: Copilot <[email protected]>

Author: wrongtest-intellif

include/tvm/tir/builtin.h

@@ -49,6 +49,14 @@ TVM_DLL const Op& ret();
  * \brief Return from a GPU thread.
  */
 TVM_DLL const Op& thread_return();
+/*!
+ * \brief Loop continue.
+ */
+TVM_DLL const Op& continue_loop();
+/*!
+ * \brief Loop break.
+ */
+TVM_DLL const Op& break_loop();
 /*!
  * \brief Reinterpret the value using the target type.
  */

include/tvm/tir/op.h

@@ -99,6 +99,20 @@ TVM_DLL PrimExpr ret(PrimExpr value, Span span = Span());
  */
 TVM_DLL PrimExpr thread_return(Span span = Span());
 
+/*!
+ * \brief Continue current loop.
+ * \param span The location of this operation in the source.
+ * \return The continue loop expression.
+ */
+TVM_DLL PrimExpr continue_loop(Span span = Span());
+
+/*!
+ * \brief Break current loop.
+ * \param span The location of this operation in the source.
+ * \return The break loop expression.
+ */
+TVM_DLL PrimExpr break_loop(Span span = Span());
+
 /*!
  * Query the maximum possible value of dtype.
  * \param dtype The data type.

include/tvm/tir/stmt.h

@@ -1310,6 +1310,9 @@ constexpr const char* explicit_read_region = "explicit_read_region";
  */
 constexpr const char* explicit_write_region = "explicit_write_region";
 
+/*! \brief ,ark a ForNode represent an irregular loop of non-structural control flow edges. */
+constexpr const char* irregular_loop_mark = "irregular_loop_mark";
+
 /*!
  * \brief Check if attr_key is a pragma key extension
  * \param attr_key The attr key to be compared

python/tvm/script/ir_builder/tir/ir.py

@@ -1917,6 +1917,8 @@ def wrapped(*args, **kwargs):
 q_multiply_shift = _op_wrapper(_tir_op.q_multiply_shift)
 q_multiply_shift_per_axis = _op_wrapper(_tir_op.q_multiply_shift_per_axis)
 ret = _op_wrapper(_tir_op.ret)
+continue_loop = _op_wrapper(_tir_op.continue_loop)
+break_loop = _op_wrapper(_tir_op.break_loop)
 round = _op_wrapper(_tir_op.round)  # pylint: disable=redefined-builtin
 rsqrt = _op_wrapper(_tir_op.rsqrt)
 shift_left = _op_wrapper(_tir_op.shift_left)
@@ -2195,6 +2197,8 @@ def wrapped(*args, **kwargs):
     "q_multiply_shift",
     "q_multiply_shift_per_axis",
     "ret",
+    "continue_loop",
+    "break_loop",
     "reinterpret",
     "round",
     "rsqrt",

python/tvm/script/parser/core/parser.py

@@ -872,6 +872,36 @@ def visit_Return(self, node: doc.Return) -> Any:  # pylint: disable=invalid-name
         """
         return _dispatch(self, "Return")(self, node)
 
+    def visit_Continue(self, node: doc.Continue) -> Any:  # pylint: disable=invalid-name
+        """The general continue visiting method.
+
+        Parameters
+        ----------
+        node : doc.Continue
+            The doc AST continue node.
+
+        Returns
+        -------
+        res : Any
+            The visiting result.
+        """
+        return _dispatch(self, "Continue")(self, node)
+
+    def visit_Break(self, node: doc.Break) -> Any:  # pylint: disable=invalid-name
+        """The general break visiting method.
+
+        Parameters
+        ----------
+        node : doc.Break
+            The doc AST break node.
+
+        Returns
+        -------
+        res : Any
+            The visiting result.
+        """
+        return _dispatch(self, "Break")(self, node)
+
     def visit_Nonlocal(self, node: doc.Nonlocal) -> Any:  # pylint: disable=invalid-name
         """The general nonlocal visiting method.
 

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

@@ -353,7 +353,8 @@ def visit_with(self: Parser, node: doc.With) -> None:
             frame = self.eval_expr(item.context_expr)
             if not isinstance(frame, Frame):
                 self.report_error(
-                    item.context_expr, "Invalid context expression in the with-statement."
+                    item.context_expr,
+                    "Invalid context expression in the with-statement.",
                 )
             rhs = stack.enter_context(frame)
             if item.optional_vars is not None:
@@ -498,7 +499,8 @@ def visit_if(self: Parser, node: doc.If) -> None:
                     self.visit_body(node.orelse)
         else:
             self.report_error(
-                node.test, f"If condition must be a boolean expression, but got {predicate}"
+                node.test,
+                f"If condition must be a boolean expression, but got {predicate}",
             )
 
 
@@ -539,6 +541,36 @@ def visit_return(self: Parser, node: doc.Return) -> None:
     T.evaluate(tvm.tir.ret(value))
 
 
+@dispatch.register(token="tir", type_name="Continue")
+def visit_continue(self: Parser, node: doc.Continue) -> None:  # pylint:disable=unused-argument
+    """The continue visiting method for tir.
+
+    Parameters
+    ----------
+    self : Parser
+        The visiting parser.
+
+    node : doc.Continue
+        The doc AST continue node.
+    """
+    T.evaluate(tvm.tir.continue_loop())
+
+
+@dispatch.register(token="tir", type_name="Break")
+def visit_break(self: Parser, node: doc.Break) -> None:  # pylint:disable=unused-argument
+    """The continue visiting method for tir.
+
+    Parameters
+    ----------
+    self : Parser
+        The visiting parser.
+
+    node : doc.Break
+        The doc AST break node.
+    """
+    T.evaluate(tvm.tir.break_loop())
+
+
 @dispatch.register(token="tir", type_name="tvm_declare_function")
 def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> GlobalVar:
     """The function declaration step for tir

python/tvm/tir/__init__.py

@@ -50,7 +50,13 @@
 from .op import tvm_stack_alloca, tvm_stack_make_shape, tvm_stack_make_array
 from .op import tvm_tuple, handle_add_byte_offset, tvm_struct_get, tvm_struct_set
 from .op import address_of, lookup_param, assume, undef
-from .op import tvm_thread_allreduce, type_annotation, tvm_access_ptr, tvm_throw_last_error
+from .op import continue_loop, break_loop
+from .op import (
+    tvm_thread_allreduce,
+    type_annotation,
+    tvm_access_ptr,
+    tvm_throw_last_error,
+)
 from .op import (
     tvm_load_matrix_sync,
     tvm_store_matrix_sync,
@@ -86,7 +92,18 @@
 from .op import tan, tanh, atan, atan2, atanh
 from .op import bitwise_and, bitwise_not, bitwise_or, bitwise_xor
 from .op import erf, sigmoid, sqrt, rsqrt, floor, ceil, hypot
-from .op import trunc, abs, round, nextafter, nearbyint, power, pow, popcount, fmod, if_then_else
+from .op import (
+    trunc,
+    abs,
+    round,
+    nextafter,
+    nearbyint,
+    power,
+    pow,
+    popcount,
+    fmod,
+    if_then_else,
+)
 from .op import likely, isnan, isnullptr, isfinite, isinf, copysign
 from .op import div, indexdiv, indexmod, truncdiv, truncmod, floordiv, floormod, ceildiv, logaddexp
 from .op import comm_reducer, min, max, sum

python/tvm/tir/op.py

@@ -1884,8 +1884,7 @@ def ret(val, span=None):
 
 
 def thread_return(span=None):
-    """Return from a GPU thread.
-
+    """Return from a GPU thread
     Parameters
     ----------
     span : Optional[Span]
@@ -1900,6 +1899,40 @@ def thread_return(span=None):
     return _ffi_api.thread_return(span)
 
 
+def continue_loop(span=None):
+    """Create a tir intrinsic call to represent continue expression
+
+    Parameters
+    ----------
+    span : Optional[Span]
+        The location of this operator in the source code.
+
+    Returns
+    -------
+    ret : PrimExpr
+        The continue expression
+    """
+
+    return _ffi_api.continue_loop(span)
+
+
+def break_loop(span=None):
+    """Create a tir intrinsic call to represent break expression
+
+    Parameters
+    ----------
+    span : Optional[Span]
+        The location of this operator in the source code.
+
+    Returns
+    -------
+    ret : PrimExpr
+        The break expression
+    """
+
+    return _ffi_api.break_loop(span)
+
+
 def any(*args, span=None):
     """Create a new experssion of the union of all conditions in the arguments
 

python/tvm/tir/pipeline.py

@@ -43,6 +43,7 @@ def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.I
             tir.transform.LowerMatchBuffer(),
             tir.transform.Simplify(),
             tir.transform.InjectPermutedLayout(),
+            tir.transform.AnnotateIrregularLoop(),
             tir.transform.InjectSoftwarePipeline(),
             tir.transform.TransformMmaBufferLayout(),
             tir.transform.LowerOpaqueBlock(),

python/tvm/tir/transform/transform.py

@@ -430,6 +430,19 @@ def AnnotateDeviceRegions():
     return _ffi_api.AnnotateDeviceRegions()  # type: ignore
 
 
+def AnnotateIrregularLoop():
+    """Annotate irregular loop mark. Loop transformations like
+    peeling, partition, unroll, etc is not allowed on irregular
+    loop with internal loop continuation and breaks.
+
+    Returns
+    -------
+    fpass : tvm.transform.Pass
+        The result pass
+    """
+    return _ffi_api.AnnotateIrregularLoop()  # type: ignore
+
+
 def SplitHostDevice():
     """Split the function into a host function and device functions.