fix ci and pass bug

Author: xwhzz

src/transform/warp_specialized_rewriter.cc

@@ -376,21 +376,32 @@ class ThreadIdxRewriter : public StmtExprMutator {
           eq_op->b.as<VarNode>() == thread_var_.get()) {
         maybe_thread_opt_ = true;
       }
-      maybe_thread_opt_ = do_shuffle_ && maybe_thread_opt_;
+      auto then_case = StmtExprMutator::VisitStmt(op->then_case);
+      maybe_thread_opt_ = do_shuffle_ && maybe_thread_opt_ && has_tma_op_;
+      if (maybe_thread_opt_) {
+        return IfThenElse(
+            Call(DataType::Bool(), tl_shuffle_elect(), {thread_extent_}),
+            StmtExprMutator::VisitStmt(op->then_case), std::nullopt);
+      }
     }
-    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);
+    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/webgpu/test_webgpu_codegen.py

@@ -43,8 +43,9 @@ def assert_gemm_codegen(
     accum_dtype="float",
 ):
     func = matmul(M, N, K, block_M, block_N, block_K, dtype=dtype, accum_dtype=accum_dtype)
-
-    artifact = tilelang.lower(func, target="webgpu")
+    # 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