[TVMScript] Fix print round-tripable multi thread env binding (#13622)

* Fix print round-tripable multi thread env binding * add unittest
apache · Dec 15, 2022 · ce97138 · ce97138
1 parent c5911a6
commit ce97138
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git a/src/printer/tvmscript_printer.cc b/src/printer/tvmscript_printer.cc
@@ -1045,7 +1045,6 @@ Doc TVMScriptPrinter::VisitStmt_(const AttrStmtNode* op) {
             << ")";
         doc << Doc::NewLine() << PrintBody(op->body);
       }
-      TryDeallocVar(iter_var->var);
       return doc;
     }
   }

diff --git a/tests/python/unittest/test_tvmscript_roundtrip.py b/tests/python/unittest/test_tvmscript_roundtrip.py
@@ -3537,6 +3537,19 @@ def func(A: T.Buffer[1, "bool"], i: T.bool, j: T.bool, k: T.bool):
         yield generator
 
 
+def multi_env_threads():
+    @T.prim_func
+    def func(A: T.Buffer[128, "float32"], C: T.Buffer[128, "float32"]):
+        B = T.alloc_buffer([128], dtype="float32")
+        for i in T.thread_binding(128, thread="threadIdx.x"):
+            B[i] = A[i] + 1.0
+        for i in T.thread_binding(128, thread="threadIdx.x"):
+            C[i] = B[i] + 2.0
+
+    mod = tvm.tir.transform.LowerOpaqueBlock()(tvm.IRModule.from_expr(func))
+    return mod["main"]
+
+
 ir_generator = tvm.testing.parameter(
     opt_gemm_normalize,
     opt_gemm_lower,
@@ -3593,6 +3606,7 @@ def func(A: T.Buffer[1, "bool"], i: T.bool, j: T.bool, k: T.bool):
     elif_chain_without_else,
     elif_chain_with_else,
     *nested_boolean_expressions(),
+    multi_env_threads,
 )