add printer/parser test, fix lint

Author: wrongtest-intellif

src/tir/op/op.cc

@@ -265,12 +265,12 @@ PrimExpr break_loop(Span span) {
 
 TVM_FFI_STATIC_INIT_BLOCK({
   namespace refl = tvm::ffi::reflection;
-  refl::GlobalDef().def("tir.thread_return", thread_return)
-    .def("tir.continue_loop", continue_loop)
-    .def("tir.break_loop", break_loop);
+  refl::GlobalDef()
+      .def("tir.thread_return", thread_return)
+      .def("tir.continue_loop", continue_loop)
+      .def("tir.break_loop", break_loop);
 });
 
-
 // maximum and min limits
 PrimExpr max_value(const DataType& dtype, Span span) {
   using namespace tir;

tests/python/tir-base/test_tir_base.py

@@ -135,11 +135,9 @@ def func(In: T.Buffer[(2,), "int32"], Out: T.Buffer[(2,), "int32"]):
     func = build_tir_func(func)
     a = np.asarray([49, 8], "int32")
     b = np.zeros([2], "int32")
-    a = tvm.nd.array(a)
-    b = tvm.nd.array(b)
     func(a, b)
-    assert b.numpy()[0] == 13
-    assert b.numpy()[1] == 9
+    assert b[0] == 13
+    assert b[1] == 9
 
 
 def test_continue_loop():
@@ -165,8 +163,8 @@ def func(Out: T.Buffer[(2,), "int32"]):
     b = np.zeros([2], "int32")
     b = tvm.nd.array(b)
     func(b)
-    assert b.numpy()[0] == 34
-    assert b.numpy()[1] == 5  # 6, 12, 18, 24, 30
+    assert b[0] == 34
+    assert b[1] == 5  # 6, 12, 18, 24, 30
 
 
 def test_exception():

tests/python/tvmscript/test_tvmscript_printer_tir.py

@@ -1046,5 +1046,34 @@ def main(A: T.Buffer((4,), "float32"), B: T.Buffer((4,), "float32")):
     _assert_print(main, expected_output)
 
 
+def test_func_with_loop_jumps():
+    from tvm.script import tir as T
+
+    @T.prim_func
+    def main(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (4,), "float32")
+        B = T.match_buffer(b, (4,), "float32")
+        for i in range(1000):
+            if i % 13 == 0:
+                A[1] = A[1] + 1
+                continue
+            if A[0] >= B[0]:
+                break
+
+    expected_output = """
+# from tvm.script import tir as T
+
+@T.prim_func
+def main(A: T.Buffer((4,), "float32"), B: T.Buffer((4,), "float32")):
+    for i in range(1000):
+        if i % 13 == 0:
+            A[1] = A[1] + T.float32(1.0)
+            T.continue_loop()
+        if A[0] >= B[0]:
+            T.break_loop()
+    """
+    _assert_print(main, expected_output)
+
+
 if __name__ == "__main__":
     tvm.testing.main()

tests/python/tvmscript/test_tvmscript_roundtrip.py

@@ -4002,6 +4002,22 @@ def func(
     return func
 
 
+def func_with_loop_jumps():
+    @T.prim_func
+    def func(In: T.Buffer((1,), "int32"), Out: T.Buffer((2,), "int32")):
+        Out[0] = 0
+        Out[1] = 0
+        for i in range(1000):
+            if i % 13 == 0:
+                Out[1] = Out[1] + 1
+                continue
+            Out[0] = Out[0] + 1
+            if Out[0] >= In[0]:
+                break
+
+    return func
+
+
 def op_of_literal():
     op_list = [
         (T.exp, 0),
@@ -4220,6 +4236,7 @@ def func(A: R.Tensor(["N"], "float16"), _: R.Prim(value="threshold")):
     return_zero_private,
     return_zero_private_with_attr,
     func_attr_with_list,
+    func_with_loop_jumps,
     *op_of_literal(),
     *relax_match_cast_struct_info_proxy(),
     relax_symbolic_size_var,