[Bugfix] Address PassContext contamination from CI and fix incorrect rewrites...

[Bugfix] Address PassContext contamination from CI and fix incorrect rewrites in warp specialized pass (#767)

* fix ci and pass bug

* fix

* try

* lint

src/transform/warp_specialized_rewriter.cc

@@ -376,21 +376,33 @@ private:
           eq_op->b.as<VarNode>() == thread_var_.get()) {
         maybe_thread_opt_ = true;
       }
-      maybe_thread_opt_ = do_shuffle_ && maybe_thread_opt_;
-    }
-    if (maybe_thread_opt_)
+      auto then_case = StmtExprMutator::VisitStmt(op->then_case);
+      maybe_thread_opt_ = do_shuffle_ && maybe_thread_opt_ && has_tma_op_;
+      has_tma_op_ = false;
+      if (maybe_thread_opt_) {
         return IfThenElse(
             Call(DataType::Bool(), tl_shuffle_elect(), {thread_extent_}),
             StmtExprMutator::VisitStmt(op->then_case), std::nullopt);
-    else
+      }
+    }
     return StmtExprMutator::VisitStmt_(op);
   }
 
+  PrimExpr VisitExpr_(const CallNode *op) final {
+    if (op->op.same_as(tl::tma_load()) ||
+        op->op.same_as(tl::tma_load_im2col()) ||
+        op->op.same_as(tl::tma_store())) {
+      has_tma_op_ = true;
+    }
+    return StmtExprMutator::VisitExpr_(op);
+  }
+
   Var thread_var_;
   PrimExpr replaced_;
   PrimExpr thread_extent_;
   bool maybe_thread_opt_ = false;
   bool do_shuffle_;
+  bool has_tma_op_ = false;
 };
 
 Block MakeGroupBlock(const Stmt &stmt,

testing/python/transform/test_tilelang_transform_lower_tile_op.py

@@ -64,6 +64,7 @@ def test_loop_tail_split(block_M, block_N, block_K, threads, vec_load_b, dtype):
                                                bx * block_N + t % (block_N // vec_load_b) *
                                                (block_N // vec_load_b) + vec], T.float16(0))
 
+    with tvm.transform.PassContext():
         mod = tvm.tir.transform.BindTarget(auto_target)(Before)
         mod = tl.transform.LowerTileOp()(mod)
         mod = tvm.tir.transform.Simplify()(mod)

testing/python/webgpu/test_webgpu_codegen.py

@@ -43,7 +43,8 @@ def assert_gemm_codegen(
     accum_dtype="float",
 ):
     func = matmul(M, N, K, block_M, block_N, block_K, dtype=dtype, accum_dtype=accum_dtype)
-
+    # Because the current pass context have been polluted by previous testing.
+    with tvm.transform.PassContext():
         artifact = tilelang.lower(func, target="webgpu")
 
     src_code = artifact.kernel_source