"Dockerfile" did not exist on "69dfbf5e976cf75e89949af4764928dde631a3f6"
warp_specialized_rewriter.cc 45.2 KB
Newer Older
1
/*!
2
 * \file warp_specialized_rewriter.cc
3
4
5
 * \brief Warp specialized Pipeline for cuda GPU (sm90+)
 */

6
#include "warp_specialized_rewriter.h"
7
8
9
10
11

namespace tvm {
namespace tl {

using namespace tir;
12
using namespace runtime;
13
using arith::IRVisitorWithAnalyzer;
14

15
16
17
18
19
20
struct LoopInfo {
  Var loop_var;
  PrimExpr extent;
  PrimExpr min;
};

21
enum class Role : uint8_t { kConsumer, kProducer, kBoth };
22

23
class ProducerBufferDetector : public StmtExprVisitor {
24
public:
25
26
  ProducerBufferDetector(
      std::unordered_set<const BufferNode *> cur_producer_buffers)
27
      : cur_producer_buffers_(std::move(cur_producer_buffers)) {}
28
29

  void clear() { has_producer_buffer_ = false; }
30
31

  void VisitExpr_(const CallNode *call) final {
32
    if (call->op.same_as(tma_load()) || call->op.same_as(tma_load_im2col())) {
33
      has_producer_buffer_ = true;
34
    }
35
    StmtExprVisitor::VisitExpr_(call);
36
37
  }

38
39
40
41
42
43
44
45
46
  void VisitExpr_(const BufferLoadNode *op) final {
    if (cur_producer_buffers_.count(op->buffer.get())) {
      has_producer_buffer_ = true;
    }
    StmtExprVisitor::VisitExpr_(op);
  }

  bool has_producer_buffer_ = false;
  std::unordered_set<const BufferNode *> cur_producer_buffers_;
47
48
49
50
};

class ProducerUsedBufferFinder : public StmtExprVisitor {
public:
51
  auto FindProducerusedBuffer(const Stmt &stmt) {
52
    producer_buffers_.clear();
53
    let_var_to_expr_.clear();
54
55
56
57
58
59
60
61
62
    std::unordered_set<const BufferNode *> last_producer_buffers_;
    for (;;) {
      VisitStmt(stmt);
      if (producer_buffers_ == last_producer_buffers_) {
        break;
      }
      last_producer_buffers_ = producer_buffers_;
    }
    return producer_buffers_;
63
64
65
66
67
68
69
  }

  void InsertBuffer(const PrimExpr &expr) {
    // Find the buffer that is used in the condition
    VarUseDefAnalyzer usage(Array<Var>{});
    usage(expr);
    for (const auto &buffer : usage.buffer_use_count_) {
70
      producer_buffers_.insert(buffer.first);
71
    }
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
    // Also collect buffers through let bindings
    CollectBuffersFromExpr(expr);
  }

  // Collect buffers from expression, following let bindings
  void CollectBuffersFromExpr(const PrimExpr &expr) {
    PostOrderVisit(expr, [this](const ObjectRef &node) {
      if (auto bl = node.as<BufferLoadNode>()) {
        producer_buffers_.insert(bl->buffer.get());
      } else if (auto var_node = node.as<VarNode>()) {
        auto var = tvm::ffi::GetRef<Var>(var_node);
        auto it = let_var_to_expr_.find(var.get());
        if (it != let_var_to_expr_.end()) {
          CollectBuffersFromExpr(it->second);
        }
      }
    });
  }

  void VisitStmt_(const LetStmtNode *op) final {
    let_var_to_expr_[op->var.get()] = op->value;
    StmtExprVisitor::VisitStmt_(op);
94
95
96
  }

  void VisitStmt_(const IfThenElseNode *op) final {
97
98
    ProducerBufferDetector producer_buffer_detector(producer_buffers_);
    producer_buffer_detector(op->then_case);
99
    if (op->else_case.defined()) {
100
      producer_buffer_detector(op->else_case.value());
101
    }
102
    if (producer_buffer_detector.has_producer_buffer_) {
103
104
105
106
107
108
      InsertBuffer(op->condition);
    }
    StmtExprVisitor::VisitStmt_(op);
  }

  void VisitStmt_(const ForNode *op) final {
109
110
111
    ProducerBufferDetector producer_buffer_detector(producer_buffers_);
    producer_buffer_detector(op->body);
    if (producer_buffer_detector.has_producer_buffer_) {
112
113
114
115
116
117
      InsertBuffer(op->min);
      InsertBuffer(op->extent);
    }
    StmtExprVisitor::VisitStmt_(op);
  }

118
119
120
121
122
123
124
  void VisitStmt_(const BufferStoreNode *op) final {
    if (producer_buffers_.count(op->buffer.get())) {
      InsertBuffer(op->value);
    }
    StmtExprVisitor::VisitStmt_(op);
  }

125
126
127
  void VisitExpr_(const CallNode *op) final {
    if (op->op.same_as(tma_load()) || op->op.same_as(tma_load_im2col())) {
      for (auto arg : op->args) {
128
129
        // Collect buffers from args, including through let bindings
        CollectBuffersFromExpr(arg);
130
131
132
133
      }
    }
  }

134
private:
135
  std::unordered_set<const BufferNode *> producer_buffers_;
136
  std::unordered_map<const VarNode *, PrimExpr> let_var_to_expr_;
137
138
};

139
class WarpSpecializedRoleMarker : public StmtVisitor {
140
public:
141
  WarpSpecializedRoleMarker(Map<Var, Buffer> buffer_data_to_buffer)
142
      : buffer_data_to_buffer_(std::move(buffer_data_to_buffer)) {}
143

144
145
  void Prepare(const Stmt &stmt) {
    ProducerUsedBufferFinder finder;
146
    producer_buffers_ = finder.FindProducerusedBuffer(stmt);
147
148
  }

149
  Role GetRole(const StmtNode *stmt) const {
150
151
152
153
154
    auto it = map_.find(stmt);
    ICHECK(it != map_.end());
    return it->second;
  }

155
  Role GetRole(const Stmt &stmt) const { return GetRole(stmt.get()); }
156

157
  void VisitStmt_(const EvaluateNode *op) final {
158
159
    Role role = Role::kConsumer;
    if (auto call = op->value.as<CallNode>()) {
160
      if (call->op.same_as(tma_load()) || call->op.same_as(tma_load_im2col())) {
161
162
163
        role = Role::kProducer;
        has_bulk_copy_ = true;
      }
164
165
166
      if (call->op.same_as(loop_break())) {
        role = Role::kBoth;
      }
167
168
169
170
    }
    SetRole(op, role);
  }

171
  void VisitStmt_(const BufferStoreNode *op) final {
172
173
    auto scope = StorageScope::Create(GetPtrStorageScope(op->buffer->data));
    bool is_shared_store = scope.rank == StorageRank::kShared;
174
    if (producer_buffers_.count(op->buffer.get())) {
175
176
177
      SetRole(op, Role::kBoth);
      return;
    }
178
179
180
181
182
183
184
    if (!is_shared_store) {
      SetRole(op, Role::kConsumer);
      return;
    }

    // Check reads from global
    Block block(/*iter_vars=*/{}, /*reads=*/{}, /*writes=*/{}, /*name_hint=*/"",
185
                /*body*/ tvm::ffi::GetRef<Stmt>(op));
186
187
188
    auto access = GetBlockReadWriteRegion(block, buffer_data_to_buffer_);
    auto reads = access[0];
    Role role = Role::kProducer;
189
190
    if (reads.empty())
      role = Role::kConsumer;
191
192
193
194
195
196
    for (auto read : reads) {
      if (read->buffer.scope() != "global") {
        role = Role::kConsumer;
        break;
      }
    }
197
198
    if (role == Role::kProducer)
      has_simt_copy_ = true;
199
200
201
    SetRole(op, role);
  }

202
  void VisitStmt_(const SeqStmtNode *op) final {
203
204
205
206
207
208
209
210
211
212
213
    StmtVisitor::VisitStmt_(op);
    auto role = GetRole(op->seq[0]);
    for (auto stmt : op->seq) {
      if (role != GetRole(stmt)) {
        role = Role::kBoth;
        break;
      }
    }
    SetRole(op, role);
  }

214
  void VisitStmt_(const IfThenElseNode *op) final {
215
216
217
218
    StmtVisitor::VisitStmt_(op);
    auto role = GetRole(op->then_case);
    if (op->else_case.defined()) {
      auto role_else = GetRole(op->else_case.value());
219
220
      if (role != role_else)
        role = Role::kBoth;
221
222
223
224
    }
    SetRole(op, role);
  }

225
  void VisitStmt_(const BlockRealizeNode *op) final {
226
227
228
229
    StmtVisitor::VisitStmt_(op);
    SetRole(op, GetRole(op->block));
  }

230
231
232
233
234
235
  void VisitStmt_(const AllocateNode *op) final {
    StmtVisitor::VisitStmt_(op);
    Role role = Role::kConsumer;
    SetRole(op, role);
  }

236
  template <class NodeType> void HandleBodyStmt(const NodeType *op) {
237
238
239
240
    StmtVisitor::VisitStmt_(op);
    SetRole(op, GetRole(op->body));
  }

241
  void VisitStmt_(const ForNode *op) final { HandleBodyStmt(op); }
242
  void VisitStmt_(const WhileNode *op) final { HandleBodyStmt(op); }
243
244
245
246
  void VisitStmt_(const LetStmtNode *op) final { HandleBodyStmt(op); }
  void VisitStmt_(const AttrStmtNode *op) final { HandleBodyStmt(op); }
  void VisitStmt_(const AssertStmtNode *op) final { HandleBodyStmt(op); }
  void VisitStmt_(const BlockNode *op) final { HandleBodyStmt(op); }
247
248
249
250
251

  bool HasProducer() { return has_simt_copy_ || has_bulk_copy_; }

  bool HasSimtCopy() { return has_simt_copy_; }

252
253
private:
  void SetRole(const StmtNode *stmt, Role role) { map_[stmt] = role; }
254
  Map<Var, Buffer> buffer_data_to_buffer_;
255
  std::unordered_map<const StmtNode *, Role> map_;
256
257
  bool has_simt_copy_ = false;
  bool has_bulk_copy_ = false;
258
  std::unordered_set<const BufferNode *> producer_buffers_;
259
260
261
};

static PrimExpr makeGetBarrier(PrimExpr barrier_id) {
262
  return Call(DataType::Handle(), get_mbarrier(), {std::move(barrier_id)});
263
264
}

265
static Stmt makeArriveBarrier(PrimExpr barrier_id, int cta_id = -1,
266
267
                              const PrimExpr &pred = 1) {
  Array<PrimExpr> args = {makeGetBarrier(std::move(barrier_id))};
268
269
270
271
272
273
  if (cta_id != -1) {
    args.push_back(cta_id);
    args.push_back(pred);
  }
  return Evaluate(
      Call(DataType::Handle(), builtin::ptx_arrive_barrier(), args));
274
275
276
}

static Stmt makeCpAsyncBarrier(PrimExpr barrier_id) {
277
  auto call = Call(DataType::Handle(), builtin::ptx_cp_async_barrier(),
278
                   {makeGetBarrier(std::move(barrier_id))});
279
280
281
282
  return Evaluate(call);
}

static Stmt makeParityWait(PrimExpr barrier_id, PrimExpr parity) {
283
  auto call = Call(DataType::Handle(), mbarrier_wait_parity(),
284
                   {makeGetBarrier(std::move(barrier_id)), std::move(parity)});
285
286
287
288
  return Evaluate(call);
}

class ProducerTraitsCollector : public StmtExprVisitor {
289
public:
290
291
  ProducerTraitsCollector() { Clear(); }

292
  void Clear() { has_simt_copy = false; }
293

294
  void Collect(const Stmt &stmt) { VisitStmt(stmt); }
295
296
297

  bool HasSimtCopy() { return has_simt_copy; }

298
private:
299
300
301
302
303
304
305
306
307
308
309
310
  void VisitStmt_(const IfThenElseNode *op) final {
    bool old_in_if_cond = in_if_cond_;
    in_if_cond_ = true;
    VisitExpr(op->condition);
    in_if_cond_ = old_in_if_cond;

    VisitStmt(op->then_case);
    if (op->else_case.defined()) {
      VisitStmt(op->else_case.value());
    }
  }

311
  void VisitExpr_(const BufferLoadNode *op) final {
312
313
314
    if (!in_if_cond_) {
      has_simt_copy = true;
    }
315
316
317
    StmtExprVisitor::VisitExpr_(op);
  }

318
  bool has_simt_copy{};
319
  bool in_if_cond_ = false;
320
321
322
323
};

// Rewrite the producer Stmt to use the correct barrier index
class MbarrierRewriter : public StmtExprMutator {
324
public:
325
326
  static Stmt Rewrite(Stmt stmt, PrimExpr barrier_id) {
    MbarrierRewriter rewriter;
327
328
    rewriter.producer_barrier_idx_ = std::move(barrier_id);
    return rewriter(std::move(stmt));
329
330
  }

331
332
private:
  PrimExpr VisitExpr_(const CallNode *op) final {
333
    auto call = Downcast<Call>(StmtExprMutator::VisitExpr_(op));
334
    if (call->op.same_as(tma_load()) || call->op.same_as(tma_load_im2col())) {
335
336
337
338
339
340
341
342
343
344
345
346
347
      auto mbar = makeGetBarrier(producer_barrier_idx_);
      auto arg0 = call->args[0].as<Call>();
      // Check if this is a 1D TMA load
      auto is_1d_tma_load =
          arg0 && !arg0.value()->op.same_as(create_tma_descriptor()) &&
          call->op.same_as(tma_load());
      if (is_1d_tma_load) {
        call.CopyOnWrite()->args.Set(2, mbar);
      } else {
        Call access_ptr = Downcast<Call>(call->args[2]);
        ICHECK(access_ptr->op.same_as(builtin::tvm_access_ptr()));
        call.CopyOnWrite()->args.Set(1, mbar);
      }
348
349
350
351
352
353
354
    }
    return call;
  }
  PrimExpr producer_barrier_idx_;
};

class ThreadIdxRewriter : public StmtExprMutator {
355
public:
356
357
358
  static Stmt Rewrite(Stmt stmt, Var thread_var, PrimExpr replaced,
                      PrimExpr thread_extent, bool do_shuffle = false) {
    auto rewriter =
359
360
361
        ThreadIdxRewriter(std::move(thread_var), std::move(replaced),
                          std::move(thread_extent), do_shuffle);
    return rewriter(std::move(stmt));
362
363
  }

364
private:
365
366
  ThreadIdxRewriter(Var thread_var, PrimExpr replaced, PrimExpr thread_extent,
                    bool do_shuffle)
367
368
      : thread_var_(std::move(thread_var)), replaced_(std::move(replaced)),
        thread_extent_(std::move(thread_extent)), do_shuffle_(do_shuffle) {}
369

370
  PrimExpr VisitExpr_(const VarNode *var) final {
371
372
373
374
375
376
377
    if (var == thread_var_.get()) {
      return replaced_;
    } else {
      return StmtExprMutator::VisitExpr_(var);
    }
  }

378
379
380
381
382
383
384
385
386
387
388
389
390
  Stmt VisitStmt_(const IfThenElseNode *op) final {
    auto f_uses_thread_index = [=](const tvm::tir::VarNode *parameter) {
      return parameter == thread_var_.get();
    };
    maybe_thread_opt_ = false;
    if (!op->else_case.defined() && op->condition.as<EQNode>() &&
        UsesVar(op->condition, f_uses_thread_index) &&
        !(UsesVar(op->then_case, f_uses_thread_index))) {
      auto eq_op = Downcast<EQ>(op->condition);
      if (eq_op->a.as<VarNode>() == thread_var_.get() ||
          eq_op->b.as<VarNode>() == thread_var_.get()) {
        maybe_thread_opt_ = true;
      }
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
      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);
      }
    }
    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;
408
    }
409
    return StmtExprMutator::VisitExpr_(op);
410
411
  }

412
413
  Var thread_var_;
  PrimExpr replaced_;
414
415
416
  PrimExpr thread_extent_;
  bool maybe_thread_opt_ = false;
  bool do_shuffle_;
417
  bool has_tma_op_ = false;
418
419
};

420
421
422
423
424
425
Block MakeGroupBlock(const Stmt &stmt,
                     const Map<String, ObjectRef> &annotations) {
  Block block(/*iter_vars=*/{}, /*reads=*/{}, /*writes=*/{}, /*name_hint=*/"",
              /*body*/ stmt,
              /*init=*/{}, /*alloc_buffers=*/{}, /*match_buffers=*/{},
              /*annotations=*/annotations);
426
427
428
429
  return block;
}

struct OpInfo {
430
  int group_size{}, order{}, stage{};
431
432
433
434
435
436
  std::vector<int> group;
};
struct PipelineInfo {
  std::vector<OpInfo> op_infos;

  PipelineInfo() = default;
437
438
439
  PipelineInfo(const Array<Array<Integer>> &group_info,
               const Array<Integer> &order_info,
               const Array<Integer> &stage_info) {
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
    int n = static_cast<int>(group_info.size());
    ICHECK(n == static_cast<int>(order_info.size()));
    ICHECK(n == static_cast<int>(stage_info.size()));
    // int cur_id = 0;
    for (int i = 0; i < n; i++) {
      OpInfo op_info;
      op_info.group_size = group_info[i].size();
      for (int j = 0; j < op_info.group_size; j++) {
        op_info.group.push_back(group_info[i][j].as<IntImmNode>()->value);
      }
      op_info.order = order_info[i].as<IntImmNode>()->value;
      op_info.stage = stage_info[i].as<IntImmNode>()->value;
      op_infos.push_back(op_info);
    }
  }

456
  PipelineInfo(const PipelineInfo &other) {
457
    for (const auto &op_info : other.op_infos) {
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
      op_infos.push_back(op_info);
    }
  }

  std::pair<int, int> FindStmt(int stmt_idx) {
    for (size_t i = 0; i < op_infos.size(); i++) {
      for (size_t j = 0; j < op_infos[i].group.size(); j++) {
        if (op_infos[i].group[j] == stmt_idx) {
          return std::make_pair(i, j);
        }
      }
    }
    return std::make_pair(-1, -1);
  }

  void UpdateOrder(int order) {
    for (int i = 0; i < static_cast<int>(op_infos.size()); i++) {
      if (op_infos[i].order >= order && op_infos[i].order > 0) {
        op_infos[i].order++;
      }
    }
  }

  int SplitOp(int stmt_idx) {
    auto pair = FindStmt(stmt_idx);
    int op_idx = pair.first;
    int inner_idx = pair.second;
    ICHECK(op_idx != -1);
    ICHECK(inner_idx != -1);
    OpInfo half0;
    OpInfo half1;
    // The order to do sync
    int sync_order = op_infos[op_idx].order + 1;
    UpdateOrder(sync_order);

    half0.group_size = inner_idx + 1;
    half0.order = op_infos[op_idx].order;
    half0.stage = op_infos[op_idx].stage;
    for (int i = 0; i <= inner_idx; i++) {
      half0.group.push_back(op_infos[op_idx].group[i]);
    }
    half1.group_size = op_infos[op_idx].group_size - inner_idx - 1;
    half1.order = op_infos[op_idx].order + 2;
    half1.stage = op_infos[op_idx].stage;
    for (int i = inner_idx + 1; i < op_infos[op_idx].group_size; i++) {
      half1.group.push_back(op_infos[op_idx].group[i]);
    }
    op_infos.erase(op_infos.begin() + op_idx);
    if (half0.group_size > 0) {
      op_infos.insert(op_infos.begin() + op_idx, half0);
    }
    if (half1.group_size > 0) {
      UpdateOrder(half1.order);
      op_infos.insert(op_infos.begin() + op_idx + 1, half1);
    }
    return sync_order;
  }

  void PrintPipelineInfo() {
517
    std::cout << "Print op_infos:" << '\n';
518
    for (size_t i = 0; i < op_infos.size(); i++) {
519
      std::cout << i << " " << op_infos[i].group_size << " "
520
                << op_infos[i].order << " " << op_infos[i].stage << '\n';
521
    }
522
    std::cout << "End of print" << '\n';
523
524
525
526
  }
};

class GroupOpRewriter : public StmtExprMutator {
527
public:
528
529
  GroupOpRewriter(const PipelineInfo &pipeline_info)
      : pipeline_info_(pipeline_info) {}
530

531
532
private:
  Stmt VisitStmt_(const ForNode *op) final {
533
534
535
536
    Map<String, ObjectRef> annotations;
    annotations.Set(String("stmt_group"), Integer(1));
    auto original_node = (op->body).as<SeqStmtNode>();
    if (!original_node) {
537
      return tvm::ffi::GetRef<For>(op);
538
539
540
541
    }
    Array<Stmt> new_body;
    int cur_id = 0;
    for (int i = 0; i < static_cast<int>(pipeline_info_.op_infos.size()); i++) {
542
543
      if (pipeline_info_.op_infos[i].group_size == 0)
        continue;
544
      Array<Stmt> block_stmt;
545
546
      for (int j = 0;
           j < static_cast<int>(pipeline_info_.op_infos[i].group_size); j++) {
547
        // ICHECK(group_info_[i][j].as<IntImmNode>());
548
549
        // int index =
        // static_cast<int>(group_info_[i][j].as<IntImmNode>()->value);
550
551
552
553
554
555
        ICHECK(original_node->seq[cur_id].as<BlockNode>());
        auto block = original_node->seq[cur_id].as<BlockNode>();
        // TODO: handle nested seqstmt
        block_stmt.push_back(block->body);
        cur_id++;
      }
556
557
558
559
560
      new_body.push_back(MakeGroupBlock(
          block_stmt.size() == 1 ? block_stmt[0]
                                 // NOLINTNEXTLINE(performance-move-const-arg)
                                 : SeqStmt(std::move(block_stmt)),
          annotations));
561
562
563
    }
    Array<Integer> order_anno;
    Array<Integer> stage_anno;
564
    for (const auto &op_info : pipeline_info_.op_infos) {
565
566
567
      order_anno.push_back(Integer(op_info.order));
      stage_anno.push_back(Integer(op_info.stage));
    }
568
    Map<String, Any> for_annotations = op->annotations;
569
570
571
    for_annotations.erase("tl_pipeline_group");
    for_annotations.Set("software_pipeline_order", order_anno);
    for_annotations.Set("software_pipeline_stage", stage_anno);
572
573
574
575
    For new_for =
        For(op->loop_var, op->min, op->extent, op->kind,
            new_body.size() == 1 ? new_body[0] : SeqStmt(std::move(new_body)),
            op->thread_binding, for_annotations);
576
577
578
579
580
    return new_for;
  }

  PipelineInfo pipeline_info_;
};
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605

class WgMMACollector : public StmtExprVisitor {
public:
  WgMMACollector() = default;

  void VisitExpr_(const CallNode *op) final {
    if (op->op.same_as(tl_gemm()) || op->op.same_as(tl_gemm_sp())) {
      auto op_name = std::string(op->args[0].as<StringImmNode>()->value);
      if (has_wgmma_) {
        has_wgmma_ =
            op_name.find("false") == std::string::npos && !in_if_scope_;
      }
    }
    StmtExprVisitor::VisitExpr_(op);
  }

  void VisitStmt_(const IfThenElseNode *op) final {
    in_if_scope_ = true;
    StmtExprVisitor::VisitStmt(op->then_case);
    if (op->else_case.defined()) {
      StmtExprVisitor::VisitStmt(op->else_case.value());
    }
    in_if_scope_ = false;
  }

606
  static bool HasWgMMA(const Stmt &stmt) {
607
608
609
610
611
612
613
614
615
    auto collector = WgMMACollector();
    collector(stmt);
    return collector.has_wgmma_;
  }

  bool has_wgmma_{true};
  bool in_if_scope_{false};
};

616
class WSCodeEmitter : public StmtMutator {
617
public:
618
  WSCodeEmitter(bool is_emitting_producer, const IterVar &thread_iv,
619
                Map<Var, Buffer> buffer_data_to_buffer,
620
                const WarpSpecializedRoleMarker &marker,
621
                bool mbarrier_only = false)
622
      : is_emitting_producer_(is_emitting_producer),
623
624
        buffer_data_to_buffer_(std::move(buffer_data_to_buffer)),
        marker_(marker), thread_var_(thread_iv->var),
625
        mbarrier_only_(mbarrier_only) {}
626

627
  /**
628
629
630
631
632
633
634
635
636
   * @brief Whether a SIMT-style bulk copy was detected.
   *
   * Returns true when a simulated SIMT (thread-parallel) copy pattern was
   * observed during analysis/emission, which can affect barrier insertion and
   * copy emission.
   *
   * @return true if a SIMT copy was detected; false otherwise.
   */
  bool hasSimtCopy() const { return has_simt_copy_; }
637

638
private:
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
  template <
      typename NodeType> /**
                          * @brief Filter a statement by its producer/consumer
                          * role for emission.
                          *
                          * Returns one of:
                          * - the original statement (unchanged) when this
                          * emitter should emit it,
                          * - the result of visiting the statement (to descend
                          * into it) when mbarrier-only mode requires full
                          * traversal for non-producer roles,
                          * - an empty evaluate (`Evaluate(0)`) when the
                          * statement should be omitted.
                          *
                          * The decision is based on the role of `op` as
                          * reported by `marker_`, the emitter mode
                          * (`is_emitting_producer_`), and the `mbarrier_only_`
                          * flag.
                          *
                          * @param op The statement node to filter; its role is
                          * queried via `marker_`.
                          * @return Stmt The statement to place into the emitted
                          * IR (possibly transformed or an empty evaluate).
                          */
  Stmt FilterByRole(const NodeType *op) {
664
    Role role = marker_.GetRole(op);
665
666
667
668
669
    if (mbarrier_only_) {
      if (role != Role::kProducer)
        return StmtMutator::VisitStmt_(op);
    }
    if (role == Role::kBoth) {
670
      return StmtMutator::VisitStmt_(op);
671
    } else if ((role == Role::kProducer) == is_emitting_producer_) {
672
      return tvm::ffi::GetRef<Stmt>(op);
673
    } else {
674
      return Evaluate(0);
675
    }
676
677
  }

678
  Stmt VisitStmt_(const SeqStmtNode *op) final {
679

680
681
682
683
684
685
686
    bool has_producer = false;
    for (auto stmt : op->seq) {
      if (marker_.GetRole(stmt) == Role::kProducer) {
        has_producer = true;
        break;
      }
    }
687
688
689
690
    bool need_producer_sync =
        has_producer && marker_.GetRole(op) == Role::kBoth;
    if (!need_producer_sync)
      return FilterByRole(op);
691

692
    auto seq_transformed =
693
        op->seq.Map([&](const Stmt &stmt) { return VisitStmt(stmt); });
694
695

    auto map = ExtractSyncPattern(op->seq);
696

697
698
699
700
701
702
703
704
705
706
707
708
709
710
    /*
      std::cout << "Print ExtractSyncPattern" << std::endl;
      for (int i = 0; i < static_cast<int>(op->seq.size()); i++) {
        std::cout << i << " " << map.acquire[i] << " " << map.release[i] << " "
        << map.release_after[i] << std::endl;
      }
      std::cout << "Print sync pattern" << std::endl;
      for (auto pattern : map.patterns) {
        std::cout << pattern.release_idx << " " << pattern.acquire_idx <<
        std::endl;
      }
      std::cout << "End of ExtractSyncPattern" << std::endl;
      pipeline_info_.PrintPipelineInfo();
    */
711
712
713
714
    Array<Stmt> new_body;
    Map<String, ObjectRef> annotations;
    annotations.Set(String("stmt_group"), Integer(1));

715
    if (is_emitting_producer_) { // producer case
716
717
718
      ProducerTraitsCollector collector;
      for (int i = 0; i < static_cast<int>(op->seq.size()); i++) {
        Array<Stmt> block_stmt = {};
719
720
721
722
723
        if (!mbarrier_only_) {
          if (marker_.GetRole(op->seq[i]) == Role::kConsumer)
            continue;
          if (marker_.GetRole(op->seq[i]) == Role::kBoth) {
            block_stmt.push_back(seq_transformed[i]);
724
725
726
727
728
729
            new_body.push_back(
                MakeGroupBlock(block_stmt.size() == 1
                                   ? block_stmt[0]
                                   // NOLINTNEXTLINE(performance-move-const-arg)
                                   : SeqStmt(std::move(block_stmt)),
                               annotations));
730
731
            continue;
          }
732
        }
733

734
        for (int pattern_idx : map.acquire[i]) {
735
          PrimExpr acquire_barrier_id =
736
737
              stage_ + num_barriers_ + num_stages_ * pattern_idx;
          PrimExpr parity = map.is_loop_dependency(pattern_idx)
738
739
                                ? bitwise_xor(parity_, 1)
                                : parity_;
740
741
          block_stmt.push_back(makeParityWait(acquire_barrier_id, parity));
        }
742
        ICHECK(!map.release[i].empty());
743
744
745
746
747
748
749
        for (size_t j = 0; j < map.release[i].size(); j++) {
          int pattern_idx = map.release[i][j];
          PrimExpr release_barrier_id =
              stage_ + num_barriers_ + num_stages_ * pattern_idx;
          auto stmt =
              MbarrierRewriter::Rewrite(seq_transformed[i], release_barrier_id);
          collector.Collect(stmt);
750
          block_stmt.push_back(stmt);
751
          if (collector.HasSimtCopy()) {
752
            block_stmt.push_back(makeCpAsyncBarrier(release_barrier_id));
753
            has_simt_copy_ = true;
754
          }
755
756
757
758
759
760
761
762
          if (map.release_after[i][j]) {
            block_stmt.push_back(makeArriveBarrier(release_barrier_id));
            for (int s = 0; s < num_stages_; s++) {
              released_barrier_.insert(s + num_barriers_ +
                                       num_stages_ * pattern_idx);
            }
          }
          collector.Clear();
763
764
765
766
767
768
          new_body.push_back(
              MakeGroupBlock(block_stmt.size() == 1
                                 ? block_stmt[0]
                                 // NOLINTNEXTLINE(performance-move-const-arg)
                                 : SeqStmt(std::move(block_stmt)),
                             annotations));
769
770
        }
      }
771
    } else { // consumer case
772
773
      for (int i = 0; i < static_cast<int>(op->seq.size()); i++) {
        Array<Stmt> block_stmt = {};
774
775
        if (marker_.GetRole(op->seq[i]) == Role::kProducer)
          continue;
776
        for (int pattern_idx : map.acquire[i]) {
777
          PrimExpr acquire_barrier_id =
778
779
              stage_ + num_barriers_ + num_stages_ * pattern_idx;
          PrimExpr parity = map.is_loop_dependency(pattern_idx)
780
781
                                ? bitwise_xor(parity_, 1)
                                : parity_;
782
783
784
          block_stmt.push_back(makeParityWait(acquire_barrier_id, parity));
        }
        block_stmt.push_back(seq_transformed[i]);
785
786
787
788
789
        for (size_t j = 0; j < map.release[i].size(); j++) {
          if (map.release_after[i][j]) {
            int pattern_idx = map.release[i][j];
            PrimExpr release_barrier_id =
                stage_ + num_barriers_ + num_stages_ * pattern_idx;
790
            block_stmt.push_back(makeArriveBarrier(release_barrier_id));
791
792
793
794
            for (int s = 0; s < num_stages_; s++) {
              released_barrier_.insert(s + num_barriers_ +
                                       num_stages_ * pattern_idx);
            }
795
796
          }
        }
797
798
799
800
801
        new_body.push_back(MakeGroupBlock(
            block_stmt.size() == 1 ? block_stmt[0]
                                   // NOLINTNEXTLINE(performance-move-const-arg)
                                   : SeqStmt(std::move(block_stmt)),
            annotations));
802
803
804
805
      }
      // Filter out the producer stmts
      int cur_id = 0;
      PipelineInfo new_pipeline_info;
806
807
      for (int i = 0; i < static_cast<int>(pipeline_info_.op_infos.size());
           i++) {
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
        auto op_info = pipeline_info_.op_infos[i];
        bool is_producer = false;
        for (int j = 0; j < op_info.group_size; j++) {
          if (marker_.GetRole(op->seq[cur_id]) == Role::kProducer) {
            is_producer = true;
          }
          cur_id++;
        }
        if (is_producer) {
          ICHECK(op_info.group_size == 1);
        } else {
          new_pipeline_info.op_infos.push_back(op_info);
        }
      }
      pipeline_info_ = new_pipeline_info;
    }

    num_barriers_ += map.patterns.size() * num_stages_;

827
    ICHECK(!new_body.empty());
828
829
830
    return new_body.size() == 1 ? new_body[0] : SeqStmt(std::move(new_body));
  }

831
  Stmt VisitStmt_(const ForNode *op) final {
832
833
    int num_stages = 1;
    auto num_stages_anno = op->annotations.Get("num_stages");
834
835
836
    if (num_stages_anno) {
      ICHECK(num_stages_anno->as<IntImmNode>());
      num_stages = static_cast<int>(num_stages_anno->as<IntImmNode>()->value);
837
838
      ICHECK(num_stages_ == 1) << "Nested pipeline not supported.";
    }
839
    loop_stack_.emplace_back(LoopInfo{op->loop_var, op->extent, op->min});
840
841
842
843

    Array<Array<Integer>> group_info_array;
    Array<Integer> order_info_array;
    Array<Integer> stage_info_array;
844

845
    auto group_anno = op->annotations.Get("tl_pipeline_group");
846
847
    if (group_anno) {
      group_info_array = Downcast<Array<Array<Integer>>>(group_anno.value());
848
849
    }
    auto order_anno = op->annotations.Get("tl_pipeline_order");
850
851
    if (order_anno) {
      order_info_array = Downcast<Array<Integer>>(order_anno.value());
852
853
    }
    auto stage_anno = op->annotations.Get("tl_pipeline_stage");
854
855
    if (stage_anno) {
      stage_info_array = Downcast<Array<Integer>>(stage_anno.value());
856
857
    }

858
859
    PipelineInfo pipeline_info(group_info_array, order_info_array,
                               stage_info_array);
860
861
    if (!pipeline_info.op_infos.empty()) {
      ICHECK(pipeline_info_.op_infos.empty())
862
          << "Nested pipeline not supported.";
863
864
865
866
867
868
869
870
871
    }

    PrimExpr parity_before = std::move(parity_);
    PrimExpr stage_before = std::move(stage_);
    int num_stages_before = num_stages_;
    PipelineInfo pipeline_info_before = pipeline_info_;

    num_stages_ = num_stages;
    pipeline_info_ = pipeline_info;
872
    PrimExpr linear_index = loop_stack_[0].loop_var - loop_stack_[0].min;
873
    for (size_t i = 1; i < loop_stack_.size(); ++i) {
874
875
      linear_index = linear_index * loop_stack_[i].extent +
                     (loop_stack_[i].loop_var - loop_stack_[i].min);
876
877
878
879
    }
    stage_ = FloorMod(linear_index, num_stages);
    parity_ = FloorMod(
        parity_before * op->extent + FloorDiv(linear_index, num_stages), 2);
880
881
882
    auto result = FilterByRole(op);

    Stmt grouped_for_node;
883
    if (result.as<ForNode>() && group_anno && !group_info_array.empty() &&
884
        !is_emitting_producer_) {
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
      GroupOpRewriter group_op_rewriter(pipeline_info_);
      auto for_node = Downcast<For>(result);
      grouped_for_node = group_op_rewriter(for_node);
    }

    parity_ = std::move(parity_before);
    stage_ = std::move(stage_before);
    num_stages_ = num_stages_before;
    pipeline_info_ = pipeline_info_before;

    // remove pipeline annotation
    auto for_node = result.as<For>();
    if (result.as<ForNode>()) {
      auto for_node = Downcast<For>(result);
      for_node.CopyOnWrite()->annotations.erase("num_stages");
900
      if (is_emitting_producer_ || group_info_array.empty()) {
901
902
903
        for_node.CopyOnWrite()->annotations.erase("tl_pipeline_order");
        for_node.CopyOnWrite()->annotations.erase("tl_pipeline_stage");
      }
904
      if (is_emitting_producer_ || !group_anno || group_info_array.empty()) {
905
        loop_stack_.pop_back();
906
907
        return for_node;
      }
908
      loop_stack_.pop_back();
909
910
      return grouped_for_node;
    }
911
    loop_stack_.pop_back();
912
913
914
    return result;
  }

915
916
917
918
919
920
  Stmt VisitStmt_(const IfThenElseNode *op) final { return FilterByRole(op); }
  Stmt VisitStmt_(const EvaluateNode *op) final { return FilterByRole(op); }
  Stmt VisitStmt_(const AttrStmtNode *op) final { return FilterByRole(op); }
  Stmt VisitStmt_(const BufferStoreNode *op) final { return FilterByRole(op); }
  Stmt VisitStmt_(const LetStmtNode *op) final { return FilterByRole(op); }
  Stmt VisitStmt_(const AssertStmtNode *op) final { return FilterByRole(op); }
921
922
  Stmt VisitStmt_(const BlockNode *op) final { return FilterByRole(op); }
  Stmt VisitStmt_(const BlockRealizeNode *op) final { return FilterByRole(op); }
923
924
925
926
927
928

  struct SyncPattern {
    int release_idx, acquire_idx;
  };

  struct SyncPatternMap {
929
930
931
    std::vector<std::vector<int>> acquire;
    std::vector<std::vector<int>> release;
    std::vector<std::vector<bool>> release_after;
932
    std::vector<SyncPattern> patterns;
933
934
935
936
937
938
939
940
941
942

    void resize(size_t n) {
      acquire.resize(n);
      release.resize(n);
      release_after.resize(n);
    }

    bool is_loop_dependency(int pattern_idx) {
      return patterns[pattern_idx].release_idx >
             patterns[pattern_idx].acquire_idx;
943
944
945
    }
  };

946
  std::vector<SyncPattern>
947
  CreateBaseSyncPairs(const Array<Stmt> &seq_stmt,
948
                      const std::vector<bool> &is_producer) {
949
    const int n = seq_stmt.size();
950
    std::vector<std::set<const BufferNode *>> reads, writes;
951
952
953
    reads.reserve(n);
    writes.reserve(n);
    for (int i = 0; i < n; i++) {
954
955
      Block block(/*iter_vars=*/{}, /*reads=*/{}, /*writes=*/{},
                  /*name_hint=*/"",
956
957
                  /*body*/ seq_stmt[i]);
      auto access = GetBlockAccessRegion(block, buffer_data_to_buffer_);
958
      std::set<const BufferNode *> read_set, write_set;
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
      for (auto region : access[0]) {
        auto var = region->buffer->data;
        if (buffer_data_to_buffer_.count(var)) {
          read_set.insert(buffer_data_to_buffer_[var].get());
        } else {
          read_set.insert(region->buffer.get());
        }
      }
      for (auto region : access[1]) {
        auto var = region->buffer->data;
        if (buffer_data_to_buffer_.count(var)) {
          write_set.insert(buffer_data_to_buffer_[var].get());
        } else {
          write_set.insert(region->buffer.get());
        }
      }
975
976
977
978
      reads.push_back(std::move(read_set));
      writes.push_back(std::move(write_set));
    }

979
980
    auto intersect_fn = [](const std::set<const BufferNode *> &lhs,
                           const std::set<const BufferNode *> &rhs) {
981
      for (auto ptr : lhs)
982
983
        if (rhs.count(ptr))
          return true;
984
985
986
987
988
989
990
991
992
      return false;
    };

    std::vector<SyncPattern> sync_patterns;
    // producer_release consumer_acquire,
    // inject before the first consumer stmt for each producer
    for (int i = 0; i < n; i++) {
      for (int j = i + 1; j < n; j++) {
        if (is_producer[i] != is_producer[j] &&
993
994
            (intersect_fn(writes[i], reads[j]) ||
             intersect_fn(reads[i], writes[j]))) {
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
          sync_patterns.push_back({i, j});
          break;
        }
      }
    }

    // consumer_release producer_acquire
    // valid when is_loop is true
    // inject before the earliest producer stmt for each consumer
    bool in_loop = !is_zero(parity_);
    if (in_loop) {
      for (int i = 0; i < n; i++) {
        for (int j = 0; j < i; j++) {
          if (is_producer[i] != is_producer[j] &&
1009
1010
              (intersect_fn(writes[i], reads[j]) ||
               intersect_fn(reads[i], writes[j]))) {
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
            sync_patterns.push_back({i, j});
            break;
          }
        }
      }
    }

    return sync_patterns;
  }

1021
1022
1023
  static std::vector<SyncPattern>
  RemoveUnusedSyncPatterns(const std::vector<SyncPattern> &sync_patterns,
                           const std::vector<bool> &is_producer) {
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
    /*
      Simplify multiple release-acquire pairs into one
      ------------------
        Produce(A)
        Produce(B)
        Consume(A, B)
      ------------------
      [(0, 2), (1, 2), (2, 0)] -> [(1, 2), (2, 0)]

      Or
      ------------------
        Produce(A, B)
        Consume(A)
        Consume(B)
      ------------------
      [(0, 1), (1, 0), (2, 0)] -> [(0, 1), (2, 0)]
    */
    int M = sync_patterns.size();
    std::vector<bool> removed(M, false);
    for (int i = 0; i < M; i++) {
      for (int j = 0; j < M; j++) {
        if (is_producer[sync_patterns[i].acquire_idx] ==
                is_producer[sync_patterns[j].acquire_idx] &&
            sync_patterns[i].acquire_idx >= sync_patterns[j].acquire_idx &&
            sync_patterns[i].release_idx < sync_patterns[j].release_idx)
          removed[i] = true;
      }
    }

    std::vector<SyncPattern> sync_pattern_cleaned;
    sync_pattern_cleaned.reserve(M);
    for (int i = 0; i < M; i++)
1056
1057
      if (!removed[i])
        sync_pattern_cleaned.push_back(sync_patterns[i]);
1058
1059
1060
1061

    return sync_pattern_cleaned;
  }

1062
  SyncPatternMap ExtractSyncPattern(const Array<Stmt> &seq_stmt) {
1063
1064
1065
1066
1067
1068
1069
1070
    size_t num_stmts = seq_stmt.size();
    std::vector<bool> is_producer;
    is_producer.reserve(num_stmts);
    for (auto stmt : seq_stmt) {
      is_producer.push_back(marker_.GetRole(stmt) == Role::kProducer);
    }

    auto sync_patterns_base = CreateBaseSyncPairs(seq_stmt, is_producer);
1071
1072
    auto sync_patterns =
        RemoveUnusedSyncPatterns(sync_patterns_base, is_producer);
1073
1074

    // for (auto pattern : sync_patterns) {
1075
1076
    //   std::cout << pattern.release_idx << " " << pattern.acquire_idx <<
    //   std::endl;
1077
1078
1079
    // }

    SyncPatternMap map;
1080
    map.resize(num_stmts);
1081
    map.patterns = sync_patterns;
1082

1083
    for (size_t i = 0; i < sync_patterns.size(); i++) {
1084
1085
1086
1087
1088
1089
      int acquire_idx = sync_patterns[i].acquire_idx;
      int release_idx = sync_patterns[i].release_idx;

      map.acquire[acquire_idx].push_back(i);
      map.release[release_idx].push_back(i);
      map.release_after[release_idx].push_back(true);
1090
1091
    }

1092
    std::vector<int> cur_consumer_barrier, cur_producer_barrier;
1093
1094
    for (int i = num_stmts - 1; i >= 0; i--) {
      if (is_producer[i]) {
1095
        if (map.release[i].empty()) {
1096
1097
1098
1099
          for (auto pattern_idx : cur_producer_barrier) {
            map.release[i].push_back(pattern_idx);
            map.release_after[i].push_back(false);
          }
1100
        } else {
1101
1102
1103
          for (auto pattern_idx : map.release[i]) {
            cur_producer_barrier.push_back(pattern_idx);
          }
1104
1105
        }
      } else {
1106
        if (map.release[i].empty()) {
1107
1108
1109
1110
          for (auto pattern_idx : cur_consumer_barrier) {
            map.release[i].push_back(pattern_idx);
            map.release_after[i].push_back(false);
          }
1111
        } else {
1112
1113
1114
          for (auto pattern_idx : map.release[i]) {
            cur_consumer_barrier.push_back(pattern_idx);
          }
1115
1116
1117
1118
1119
1120
1121
1122
1123
        }
      }
    }
    return map;
  }

  const bool is_emitting_producer_;
  Map<Var, Buffer> buffer_data_to_buffer_;
  std::unordered_set<int> released_barrier_;
1124
  const WarpSpecializedRoleMarker &marker_;
1125
1126
1127
1128
1129

  int num_barriers_ = 0;
  PrimExpr parity_ = 0;
  PrimExpr stage_ = 0;
  int num_stages_ = 1;
1130
  std::vector<LoopInfo> loop_stack_;
1131
  Var thread_var_;
1132
  bool mbarrier_only_ = false;
1133
1134
  PipelineInfo pipeline_info_;
  friend class WarpSpecializedRewriter;
1135
  bool has_simt_copy_ = false;
1136
1137
1138
};

class WarpSpecializedRewriter : public StmtExprMutator {
1139
public:
1140
1141
1142
1143
1144
1145
  WarpSpecializedRewriter(bool disable_warp_specialized,
                          bool disable_shuffle_elect)
      : disable_warp_specialized_(disable_warp_specialized),
        disable_shuffle_elect_(disable_shuffle_elect) {}
  static PrimFunc Substitute(PrimFunc f, bool disable_warp_specialized,
                             bool disable_shuffle_elect) {
1146
1147
1148
    // Check if function only uses threadIdx.x before proceeding
    if (!ThreadTagChecker::HasOnlyThreadIdxX(f)) {
      LOG(WARNING) << "WarpSpecialize will be disabled because the program "
1149
                      "uses thread tags other than threadIdx.x."
1150
1151
1152
1153
1154
1155
                   << "If you want to use warp specialization, please refactor "
                      "your program to use threadIdx.x only";
      // Return original function unchanged if other thread tags are found
      return f;
    }

1156
1157
    auto T = WarpSpecializedRewriter(disable_warp_specialized,
                                     disable_shuffle_elect);
1158
    T.buffer_lca_ = DetectBufferAccessLCA(f);
1159
1160
    for (auto [buffer, _] : T.buffer_lca_)
      T.buffer_data_to_buffer_.Set(buffer->data, buffer);
1161
1162
1163
1164
    f.CopyOnWrite()->body = T(f->body);
    return f;
  }

1165
1166
private:
  Stmt VisitStmt_(const AttrStmtNode *op) final {
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
    if (op->attr_key == tir::attr::thread_extent &&
        Downcast<IterVar>(op->node)->thread_tag == "threadIdx.x") {
      thread_iv_ = Downcast<IterVar>(op->node);
      need_update_thread_extent_ = false;
      AttrStmt attr_stmt = Downcast<AttrStmt>(StmtExprMutator::VisitStmt_(op));
      if (need_update_thread_extent_) {
        thread_iv_.CopyOnWrite()->dom = {0, updated_thread_extent_.value()};
        attr_stmt.CopyOnWrite()->node = thread_iv_;
        attr_stmt.CopyOnWrite()->value = updated_thread_extent_.value();
      }
      thread_iv_ = {};
      return attr_stmt;
    } else {
      return StmtExprMutator::VisitStmt_(op);
    }
  }

1184
1185
1186
1187
  // If users define a thread binding, we will replace the thread binding with
  // threadIdx.x We require the thread binding is threadIdx.x, and the extent is
  // the same as the thread extent
  Stmt VisitStmt_(const ForNode *op) final {
1188
1189
1190
1191
1192
1193
1194
    ICHECK(thread_iv_.defined());
    For for_node = Downcast<For>(StmtExprMutator::VisitStmt_(op));
    if (for_node->kind == ForKind::kThreadBinding) {
      ICHECK(for_node->thread_binding.defined());
      String thread_tag = for_node->thread_binding.value()->thread_tag;
      ICHECK(thread_tag == "threadIdx.x") << "Only support threadIdx.x";
      Var thread_iv = Downcast<Var>(for_node->loop_var);
1195
      Stmt new_body =
1196
          ThreadIdxRewriter::Rewrite(for_node->body, thread_iv, thread_iv_, 0);
1197
1198
1199
1200
1201
      return new_body;
    }
    return for_node;
  }

1202
1203
1204
  Stmt VisitStmt_(const BlockRealizeNode *op) final {
    BlockRealize block_realize =
        Downcast<BlockRealize>(StmtExprMutator::VisitStmt_(op));
1205
1206
1207
1208
1209
1210
    if (!thread_iv_.defined()) {
      return block_realize;
    }

    Block block = block_realize->block;
    WarpSpecializedRoleMarker marker(buffer_data_to_buffer_);
1211
    marker.Prepare(block);
1212
1213
1214
1215
1216
1217
    marker(block);
    if (!marker.HasProducer()) {
      // Cannot detect any producer here, directly return.
      return block_realize;
    }

1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
    if (disable_warp_specialized_) {
      WSCodeEmitter mbarrier_emitter(true, thread_iv_, buffer_data_to_buffer_,
                                     marker, true);
      auto code = mbarrier_emitter(block->body);
      int num_barriers = mbarrier_emitter.num_barriers_;
      Array<PrimExpr> barrier_num_threads;
      barrier_num_threads.reserve(num_barriers);
      PrimExpr arrive_thread_count = thread_iv_->dom->extent;
      for (int i = 0; i < num_barriers; i++) {
        barrier_num_threads.push_back(arrive_thread_count);
      }
      Stmt init_barrier = Evaluate(Call(
1230
          DataType::Handle(), create_list_of_mbarrier(), barrier_num_threads));
1231
1232
1233
1234
      block.CopyOnWrite()->body = SeqStmt({init_barrier, code});
      block_realize.CopyOnWrite()->block = block;
      return block_realize;
    }
1235
    WSCodeEmitter producer(true, thread_iv_, buffer_data_to_buffer_, marker);
1236
    WSCodeEmitter consumer(false, thread_iv_, buffer_data_to_buffer_, marker,
1237
                           false);
1238
1239
1240
1241
1242
    Stmt producer_code = producer(block->body);
    Stmt consumer_code = consumer(block->body);
    PrimExpr consumer_thread_extent = thread_iv_->dom->extent;
    PrimExpr producer_thread_extent = thread_iv_->dom->extent;
    // Need one warp-group for bulk-copy only case
1243
1244
    if (!marker.HasSimtCopy())
      producer_thread_extent = 128;
1245
1246

    updated_thread_extent_ = consumer_thread_extent + producer_thread_extent;
1247
1248
1249
1250
1251
1252
1253
1254

    producer_code = ThreadIdxRewriter::Rewrite(
        producer_code, thread_iv_->var,
        thread_iv_->var - consumer_thread_extent, producer_thread_extent,
        !disable_shuffle_elect_);
    consumer_code = ThreadIdxRewriter::Rewrite(
        consumer_code, thread_iv_->var, thread_iv_->var, consumer_thread_extent,
        !disable_shuffle_elect_);
1255
1256
1257
1258
1259
1260
1261
1262
    need_update_thread_extent_ = true;

    ICHECK(producer.num_barriers_ == consumer.num_barriers_)
        << producer.num_barriers_ << " " << consumer.num_barriers_;
    int num_barriers = consumer.num_barriers_;
    Array<PrimExpr> barrier_num_threads;
    barrier_num_threads.reserve(num_barriers);
    for (int i = 0; i < num_barriers; i++) {
1263
1264
1265
      PrimExpr arrive_thread_count =
          producer.released_barrier_.count(i)
              ? (producer.hasSimtCopy() ? producer_thread_extent : 1)
1266
              : consumer_thread_extent;
1267
1268
1269
      barrier_num_threads.push_back(arrive_thread_count);
    }

1270
    Stmt init_barrier = Evaluate(Call(
1271
        DataType::Handle(), create_list_of_mbarrier(), barrier_num_threads));
1272
1273
    Stmt body = IfThenElse(GE(thread_iv_->var, consumer_thread_extent),
                           producer_code, consumer_code);
1274
    // Add an attr here to handle the partial thread count in ThreadSync pass.
1275
1276
    Array<IntImm> ws_partition = {Downcast<IntImm>(producer_thread_extent),
                                  Downcast<IntImm>(consumer_thread_extent)};
1277
    body = AttrStmt(ws_partition, attr::kWarpSpecializationScope, 0, body);
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291

    block.CopyOnWrite()->body = SeqStmt({init_barrier, body});
    block_realize.CopyOnWrite()->block = block;
    return block_realize;
  }

  WarpSpecializedRewriter() = default;

  Map<Var, Buffer> buffer_data_to_buffer_;
  Map<Buffer, Optional<Stmt>> buffer_lca_;
  Map<Buffer, Buffer> buffer_remap_;
  IterVar thread_iv_;
  Optional<PrimExpr> updated_thread_extent_;
  bool need_update_thread_extent_ = false;
1292
  bool disable_warp_specialized_ = false;
1293
  bool disable_shuffle_elect_ = false;
1294
1295
1296
1297
1298
};

using namespace tir::transform;

tvm::transform::Pass WarpSpecialized() {
1299
  auto pass_func = [=](PrimFunc f, const IRModule &m, PassContext ctx) {
1300
1301
    bool disable_warp_specialized =
        ctx->GetConfig<Bool>(kDisableWarpSpecialized, Bool(false)).value();
1302
1303
    bool disable_shuffle_elect =
        ctx->GetConfig<Bool>(kDisableShuffleElect, Bool(false)).value();
1304
1305
1306
    bool warp_specialized = WarpSpecializedDetector::Detect(f->body);

    if (!warp_specialized) {
1307
1308
      return WarpSpecializedRewriter::Substitute(f, disable_warp_specialized,
                                                 disable_shuffle_elect);
1309
    } else {
1310
      auto node = ffi::String("default");
1311
1312
1313
      f.CopyOnWrite()->body =
          AttrStmt(node, attr::kCustomWarpSpecialization, 1, f->body);
      return f;
1314
    }
1315
1316
1317
1318
  };
  return CreatePrimFuncPass(pass_func, 0, "tl.WarpSpecialized", {});
}

1319
TVM_FFI_STATIC_INIT_BLOCK() {
1320
1321
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def("tl.transform.WarpSpecialized", WarpSpecialized);
1322
}
1323

1324
1325
} // namespace tl
} // namespace tvm