[Bugfix] Disable SingleEnvThreadVerifier (#16361)

During TensorIR scheduling, the `IterVar`s that represent environment
threads may duplicate, i.e. it is legal to have two env threads with
the same name tag, which may fail the `SingleEnvThreadVerifier` check
during schedule creation. This PR disables this check in this case.
In the future, it may be worthwhile to bring it back against
post-scheduling TIR.

Author: junrushao

src/tir/analysis/verify_well_formed.cc

@@ -347,7 +347,6 @@ bool VerifyWellFormed(const PrimFunc& func, bool assert_mode) {
   }
 
   if (!UndefinedVarVerifier::Verify(func, assert_mode)) return false;
-  if (!SingleEnvThreadVerifier::Verify(func, assert_mode)) return false;
 
   // TODO(Siyuan): add more checks here.
   return true;
@@ -364,7 +363,6 @@ bool VerifyWellFormed(const IRModule& mod, bool assert_mode) {
   }
 
   if (!UndefinedVarVerifier::Verify(mod, assert_mode)) return false;
-  if (!SingleEnvThreadVerifier::Verify(mod, assert_mode)) return false;
 
   return true;
 }

tests/python/tir-schedule/test_tir_schedule_error.py

@@ -15,9 +15,8 @@
 # specific language governing permissions and limitations
 # under the License.
 # pylint: disable=missing-function-docstring,missing-module-docstring
-import sys
-
 import pytest
+
 import tvm
 import tvm.testing
 from tvm import tir
@@ -41,6 +40,25 @@ def matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
                 C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
 
 
+@T.prim_func
+def two_kernels(var_A: T.handle, var_B: T.handle, seq_len: T.int32):
+    T.func_attr({"tir.noalias": T.bool(True)})
+    A = T.match_buffer(var_A, (1, seq_len * 8), "int32")
+    B = T.match_buffer(var_B, (1, seq_len * 8), "int32", align=8)
+    with T.block("exclusive_scan"):
+        T.reads()
+        T.writes()
+        s8: T.int32 = seq_len * 8
+        if s8 == 0:
+            blockIdx_x = T.launch_thread("blockIdx.x", 1)
+        else:
+            with T.launch_thread("threadIdx.x", 1024) as threadIdx_x:
+                blockIdx_x = T.launch_thread("blockIdx.x", T.ceildiv(s8, 1024))
+                i: T.int32 = blockIdx_x * 1024 + threadIdx_x
+                if i < s8:
+                    B[i // s8, i % s8] = A[i // s8, i % s8]
+
+
 # pylint: enable=no-member,invalid-name,unused-variable
 
 
@@ -74,5 +92,9 @@ def test_tir_schedule_attribute_error():
         sch.non_existent_field()
 
 
+def test_tir_schedule_two_kernels():
+    tir.Schedule(two_kernels)
+
+
 if __name__ == "__main__":
     tvm.testing.main()