warp_specialized_rewriter.cc 46.4 KB
Newer Older
1
/*!
2
 * \file warp_specialized_rewriter.cc
3
4
5
 * \brief Warp specialized Pipeline for cuda GPU (sm90+)
 */

6
#include "arith/ir_visitor_with_analyzer.h"
7
#include "tir/analysis/var_use_def_analysis.h"
8
#include <tvm/ffi/reflection/registry.h>
9
10
11
12
13
14
#include <tvm/tir/analysis.h>
#include <tvm/tir/builtin.h>
#include <tvm/tir/op.h>
#include <tvm/tir/stmt_functor.h>
#include <tvm/tir/transform.h>

15
16
#include <utility>

17
#include "../op/builtin.h"
18
#include "./common/collector.h"
19
20
#include "runtime/thread_storage_scope.h"
#include "tir/transforms/ir_utils.h"
21
22
23
24
25

namespace tvm {
namespace tl {

using namespace tir;
26
using namespace runtime;
27
using arith::IRVisitorWithAnalyzer;
28

29
30
31
32
33
34
struct LoopInfo {
  Var loop_var;
  PrimExpr extent;
  PrimExpr min;
};

35
enum class Role : uint8_t { kConsumer, kProducer, kBoth };
36

37
class ProducerBufferDetector : public StmtExprVisitor {
38
public:
39
40
  ProducerBufferDetector(
      std::unordered_set<const BufferNode *> cur_producer_buffers)
41
      : cur_producer_buffers_(std::move(cur_producer_buffers)) {}
42
43

  void clear() { has_producer_buffer_ = false; }
44
45

  void VisitExpr_(const CallNode *call) final {
46
    if (call->op.same_as(tma_load()) || call->op.same_as(tma_load_im2col())) {
47
      has_producer_buffer_ = true;
48
    }
49
    StmtExprVisitor::VisitExpr_(call);
50
51
  }

52
53
54
55
56
57
58
59
60
  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_;
61
62
63
64
};

class ProducerUsedBufferFinder : public StmtExprVisitor {
public:
65
  auto FindProducerusedBuffer(const Stmt &stmt) {
66
67
68
69
70
71
72
73
74
75
    producer_buffers_.clear();
    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_;
76
77
78
79
80
81
82
  }

  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_) {
83
      producer_buffers_.insert(buffer.first);
84
85
86
87
    }
  }

  void VisitStmt_(const IfThenElseNode *op) final {
88
89
    ProducerBufferDetector producer_buffer_detector(producer_buffers_);
    producer_buffer_detector(op->then_case);
90
    if (op->else_case.defined()) {
91
      producer_buffer_detector(op->else_case.value());
92
    }
93
    if (producer_buffer_detector.has_producer_buffer_) {
94
95
96
97
98
99
      InsertBuffer(op->condition);
    }
    StmtExprVisitor::VisitStmt_(op);
  }

  void VisitStmt_(const ForNode *op) final {
100
101
102
    ProducerBufferDetector producer_buffer_detector(producer_buffers_);
    producer_buffer_detector(op->body);
    if (producer_buffer_detector.has_producer_buffer_) {
103
104
105
106
107
108
      InsertBuffer(op->min);
      InsertBuffer(op->extent);
    }
    StmtExprVisitor::VisitStmt_(op);
  }

109
110
111
112
113
114
115
  void VisitStmt_(const BufferStoreNode *op) final {
    if (producer_buffers_.count(op->buffer.get())) {
      InsertBuffer(op->value);
    }
    StmtExprVisitor::VisitStmt_(op);
  }

116
117
118
119
  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) {
        if (auto buffer_load = arg.as<BufferLoadNode>()) {
120
          producer_buffers_.insert(buffer_load->buffer.get());
121
122
123
124
125
        }
      }
    }
  }

126
private:
127
  std::unordered_set<const BufferNode *> producer_buffers_;
128
129
};

130
class WarpSpecializedRoleMarker : public StmtVisitor {
131
public:
132
  WarpSpecializedRoleMarker(Map<Var, Buffer> buffer_data_to_buffer)
133
      : buffer_data_to_buffer_(std::move(buffer_data_to_buffer)) {}
134

135
136
  void Prepare(const Stmt &stmt) {
    ProducerUsedBufferFinder finder;
137
    producer_buffers_ = finder.FindProducerusedBuffer(stmt);
138
139
  }

140
  Role GetRole(const StmtNode *stmt) const {
141
142
143
144
145
    auto it = map_.find(stmt);
    ICHECK(it != map_.end());
    return it->second;
  }

146
  Role GetRole(const Stmt &stmt) const { return GetRole(stmt.get()); }
147

148
  void VisitStmt_(const EvaluateNode *op) final {
149
150
    Role role = Role::kConsumer;
    if (auto call = op->value.as<CallNode>()) {
151
      if (call->op.same_as(tma_load()) || call->op.same_as(tma_load_im2col())) {
152
153
154
        role = Role::kProducer;
        has_bulk_copy_ = true;
      }
155
156
157
      if (call->op.same_as(loop_break())) {
        role = Role::kBoth;
      }
158
159
160
161
    }
    SetRole(op, role);
  }

162
  void VisitStmt_(const BufferStoreNode *op) final {
163
164
    auto scope = StorageScope::Create(GetPtrStorageScope(op->buffer->data));
    bool is_shared_store = scope.rank == StorageRank::kShared;
165
    if (producer_buffers_.count(op->buffer.get())) {
166
167
168
      SetRole(op, Role::kBoth);
      return;
    }
169
170
171
172
173
174
175
176
177
178
179
    if (!is_shared_store) {
      SetRole(op, Role::kConsumer);
      return;
    }

    // Check reads from global
    Block block(/*iter_vars=*/{}, /*reads=*/{}, /*writes=*/{}, /*name_hint=*/"",
                /*body*/ GetRef<Stmt>(op));
    auto access = GetBlockReadWriteRegion(block, buffer_data_to_buffer_);
    auto reads = access[0];
    Role role = Role::kProducer;
180
181
    if (reads.empty())
      role = Role::kConsumer;
182
183
184
185
186
187
    for (auto read : reads) {
      if (read->buffer.scope() != "global") {
        role = Role::kConsumer;
        break;
      }
    }
188
189
    if (role == Role::kProducer)
      has_simt_copy_ = true;
190
191
192
    SetRole(op, role);
  }

193
  void VisitStmt_(const SeqStmtNode *op) final {
194
195
196
197
198
199
200
201
202
203
204
    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);
  }

205
  void VisitStmt_(const IfThenElseNode *op) final {
206
207
208
209
    StmtVisitor::VisitStmt_(op);
    auto role = GetRole(op->then_case);
    if (op->else_case.defined()) {
      auto role_else = GetRole(op->else_case.value());
210
211
      if (role != role_else)
        role = Role::kBoth;
212
213
214
215
    }
    SetRole(op, role);
  }

216
  void VisitStmt_(const BlockRealizeNode *op) final {
217
218
219
220
    StmtVisitor::VisitStmt_(op);
    SetRole(op, GetRole(op->block));
  }

221
222
223
224
225
226
  void VisitStmt_(const AllocateNode *op) final {
    StmtVisitor::VisitStmt_(op);
    Role role = Role::kConsumer;
    SetRole(op, role);
  }

227
  template <class NodeType> void HandleBodyStmt(const NodeType *op) {
228
229
230
231
    StmtVisitor::VisitStmt_(op);
    SetRole(op, GetRole(op->body));
  }

232
  void VisitStmt_(const ForNode *op) final { HandleBodyStmt(op); }
233
  void VisitStmt_(const WhileNode *op) final { HandleBodyStmt(op); }
234
235
236
237
  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); }
238
239
240
241
242

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

  bool HasSimtCopy() { return has_simt_copy_; }

243
244
private:
  void SetRole(const StmtNode *stmt, Role role) { map_[stmt] = role; }
245
  Map<Var, Buffer> buffer_data_to_buffer_;
246
  std::unordered_map<const StmtNode *, Role> map_;
247
248
  bool has_simt_copy_ = false;
  bool has_bulk_copy_ = false;
249
  std::unordered_set<const BufferNode *> producer_buffers_;
250
251
252
};

static PrimExpr makeGetBarrier(PrimExpr barrier_id) {
253
  return Call(DataType::Handle(), get_mbarrier(), {std::move(barrier_id)});
254
255
}

256
static Stmt makeArriveBarrier(PrimExpr barrier_id, int cta_id = -1,
257
258
                              const PrimExpr &pred = 1) {
  Array<PrimExpr> args = {makeGetBarrier(std::move(barrier_id))};
259
260
261
262
263
264
  if (cta_id != -1) {
    args.push_back(cta_id);
    args.push_back(pred);
  }
  return Evaluate(
      Call(DataType::Handle(), builtin::ptx_arrive_barrier(), args));
265
266
267
}

static Stmt makeCpAsyncBarrier(PrimExpr barrier_id) {
268
  auto call = Call(DataType::Handle(), builtin::ptx_cp_async_barrier(),
269
                   {makeGetBarrier(std::move(barrier_id))});
270
271
272
273
  return Evaluate(call);
}

static Stmt makeParityWait(PrimExpr barrier_id, PrimExpr parity) {
274
  auto call = Call(DataType::Handle(), mbarrier_wait_parity(),
275
                   {makeGetBarrier(std::move(barrier_id)), std::move(parity)});
276
277
278
279
  return Evaluate(call);
}

class ProducerTraitsCollector : public StmtExprVisitor {
280
public:
281
282
  ProducerTraitsCollector() { Clear(); }

283
  void Clear() { has_simt_copy = false; }
284

285
  void Collect(const Stmt &stmt) { VisitStmt(stmt); }
286
287
288

  bool HasSimtCopy() { return has_simt_copy; }

289
private:
290
291
292
293
294
295
296
297
298
299
300
301
  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());
    }
  }

302
  void VisitExpr_(const BufferLoadNode *op) final {
303
304
305
    if (!in_if_cond_) {
      has_simt_copy = true;
    }
306
307
308
    StmtExprVisitor::VisitExpr_(op);
  }

309
  bool has_simt_copy{};
310
  bool in_if_cond_ = false;
311
312
313
314
};

// Rewrite the producer Stmt to use the correct barrier index
class MbarrierRewriter : public StmtExprMutator {
315
public:
316
317
  static Stmt Rewrite(Stmt stmt, PrimExpr barrier_id) {
    MbarrierRewriter rewriter;
318
319
    rewriter.producer_barrier_idx_ = std::move(barrier_id);
    return rewriter(std::move(stmt));
320
321
  }

322
323
private:
  PrimExpr VisitExpr_(const CallNode *op) final {
324
    auto call = Downcast<Call>(StmtExprMutator::VisitExpr_(op));
325
    if (call->op.same_as(tma_load()) || call->op.same_as(tma_load_im2col())) {
326
327
328
329
330
331
332
333
334
335
336
337
338
      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);
      }
339
340
341
342
343
344
345
    }
    return call;
  }
  PrimExpr producer_barrier_idx_;
};

class ThreadIdxRewriter : public StmtExprMutator {
346
public:
347
348
349
  static Stmt Rewrite(Stmt stmt, Var thread_var, PrimExpr replaced,
                      PrimExpr thread_extent, bool do_shuffle = false) {
    auto rewriter =
350
351
352
        ThreadIdxRewriter(std::move(thread_var), std::move(replaced),
                          std::move(thread_extent), do_shuffle);
    return rewriter(std::move(stmt));
353
354
  }

355
private:
356
357
  ThreadIdxRewriter(Var thread_var, PrimExpr replaced, PrimExpr thread_extent,
                    bool do_shuffle)
358
359
      : thread_var_(std::move(thread_var)), replaced_(std::move(replaced)),
        thread_extent_(std::move(thread_extent)), do_shuffle_(do_shuffle) {}
360

361
  PrimExpr VisitExpr_(const VarNode *var) final {
362
363
364
365
366
367
368
    if (var == thread_var_.get()) {
      return replaced_;
    } else {
      return StmtExprMutator::VisitExpr_(var);
    }
  }

369
370
371
372
373
374
375
376
377
378
379
380
381
  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;
      }
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
      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;
399
    }
400
    return StmtExprMutator::VisitExpr_(op);
401
402
  }

403
404
  Var thread_var_;
  PrimExpr replaced_;
405
406
407
  PrimExpr thread_extent_;
  bool maybe_thread_opt_ = false;
  bool do_shuffle_;
408
  bool has_tma_op_ = false;
409
410
};

411
412
413
414
415
416
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);
417
418
419
420
  return block;
}

struct OpInfo {
421
  int group_size{}, order{}, stage{};
422
423
424
425
426
427
  std::vector<int> group;
};
struct PipelineInfo {
  std::vector<OpInfo> op_infos;

  PipelineInfo() = default;
428
429
430
  PipelineInfo(const Array<Array<Integer>> &group_info,
               const Array<Integer> &order_info,
               const Array<Integer> &stage_info) {
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
    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);
    }
  }

447
  PipelineInfo(const PipelineInfo &other) {
448
    for (const auto &op_info : other.op_infos) {
449
450
451
452
453
454
455
456
457
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
      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() {
508
    std::cout << "Print op_infos:" << '\n';
509
    for (size_t i = 0; i < op_infos.size(); i++) {
510
      std::cout << i << " " << op_infos[i].group_size << " "
511
                << op_infos[i].order << " " << op_infos[i].stage << '\n';
512
    }
513
    std::cout << "End of print" << '\n';
514
515
516
517
  }
};

class GroupOpRewriter : public StmtExprMutator {
518
public:
519
520
  GroupOpRewriter(const PipelineInfo &pipeline_info)
      : pipeline_info_(pipeline_info) {}
521

522
523
private:
  Stmt VisitStmt_(const ForNode *op) final {
524
525
526
527
528
529
530
531
532
    Map<String, ObjectRef> annotations;
    annotations.Set(String("stmt_group"), Integer(1));
    auto original_node = (op->body).as<SeqStmtNode>();
    if (!original_node) {
      return GetRef<For>(op);
    }
    Array<Stmt> new_body;
    int cur_id = 0;
    for (int i = 0; i < static_cast<int>(pipeline_info_.op_infos.size()); i++) {
533
534
      if (pipeline_info_.op_infos[i].group_size == 0)
        continue;
535
      Array<Stmt> block_stmt;
536
537
      for (int j = 0;
           j < static_cast<int>(pipeline_info_.op_infos[i].group_size); j++) {
538
        // ICHECK(group_info_[i][j].as<IntImmNode>());
539
540
        // int index =
        // static_cast<int>(group_info_[i][j].as<IntImmNode>()->value);
541
542
543
544
545
546
        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++;
      }
547
548
549
550
      new_body.push_back(MakeGroupBlock(block_stmt.size() == 1
                                            ? block_stmt[0]
                                            : SeqStmt(std::move(block_stmt)),
                                        annotations));
551
552
553
    }
    Array<Integer> order_anno;
    Array<Integer> stage_anno;
554
    for (const auto &op_info : pipeline_info_.op_infos) {
555
556
557
      order_anno.push_back(Integer(op_info.order));
      stage_anno.push_back(Integer(op_info.stage));
    }
558
    Map<String, Any> for_annotations = op->annotations;
559
560
561
    for_annotations.erase("tl_pipeline_group");
    for_annotations.Set("software_pipeline_order", order_anno);
    for_annotations.Set("software_pipeline_stage", stage_anno);
562
563
564
565
    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);
566
567
568
569
570
    return new_for;
  }

  PipelineInfo pipeline_info_;
};
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595

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;
  }

596
  static bool HasWgMMA(const Stmt &stmt) {
597
598
599
600
601
602
603
604
605
    auto collector = WgMMACollector();
    collector(stmt);
    return collector.has_wgmma_;
  }

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

606
class WSCodeEmitter : public StmtMutator {
607
public:
608
  WSCodeEmitter(bool is_emitting_producer, const IterVar &thread_iv,
609
                Map<Var, Buffer> buffer_data_to_buffer,
610
                const WarpSpecializedRoleMarker &marker,
611
                bool mbarrier_only = false)
612
      : is_emitting_producer_(is_emitting_producer),
613
614
        buffer_data_to_buffer_(std::move(buffer_data_to_buffer)),
        marker_(marker), thread_var_(thread_iv->var),
615
        mbarrier_only_(mbarrier_only) {}
616

617
  /**
618
619
620
621
622
623
624
625
626
   * @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_; }
627

628
private:
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
  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) {
654
    Role role = marker_.GetRole(op);
655
656
657
658
659
    if (mbarrier_only_) {
      if (role != Role::kProducer)
        return StmtMutator::VisitStmt_(op);
    }
    if (role == Role::kBoth) {
660
      return StmtMutator::VisitStmt_(op);
661
    } else if ((role == Role::kProducer) == is_emitting_producer_) {
662
      return GetRef<Stmt>(op);
663
    } else {
664
      return Evaluate(0);
665
    }
666
667
  }

668
  Stmt VisitStmt_(const SeqStmtNode *op) final {
669

670
671
672
673
674
675
676
    bool has_producer = false;
    for (auto stmt : op->seq) {
      if (marker_.GetRole(stmt) == Role::kProducer) {
        has_producer = true;
        break;
      }
    }
677
678
679
680
    bool need_producer_sync =
        has_producer && marker_.GetRole(op) == Role::kBoth;
    if (!need_producer_sync)
      return FilterByRole(op);
681

682
    auto seq_transformed =
683
        op->seq.Map([&](const Stmt &stmt) { return VisitStmt(stmt); });
684
685

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

687
688
689
690
691
692
693
694
695
696
697
698
699
700
    /*
      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();
    */
701
702
703
704
    Array<Stmt> new_body;
    Map<String, ObjectRef> annotations;
    annotations.Set(String("stmt_group"), Integer(1));

705
    if (is_emitting_producer_) { // producer case
706
707
708
      ProducerTraitsCollector collector;
      for (int i = 0; i < static_cast<int>(op->seq.size()); i++) {
        Array<Stmt> block_stmt = {};
709
710
711
712
713
714
715
716
717
718
719
        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]);
            new_body.push_back(MakeGroupBlock(
                block_stmt.size() == 1 ? block_stmt[0]
                                       : SeqStmt(std::move(block_stmt)),
                annotations));
            continue;
          }
720
        }
721

722
        for (int pattern_idx : map.acquire[i]) {
723
          PrimExpr acquire_barrier_id =
724
725
              stage_ + num_barriers_ + num_stages_ * pattern_idx;
          PrimExpr parity = map.is_loop_dependency(pattern_idx)
726
727
                                ? bitwise_xor(parity_, 1)
                                : parity_;
728
729
          block_stmt.push_back(makeParityWait(acquire_barrier_id, parity));
        }
730
        ICHECK(!map.release[i].empty());
731
732
733
734
735
736
737
        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);
738
          block_stmt.push_back(stmt);
739
          if (collector.HasSimtCopy()) {
740
            block_stmt.push_back(makeCpAsyncBarrier(release_barrier_id));
741
            has_simt_copy_ = true;
742
          }
743
744
745
746
747
748
749
750
751
752
753
754
          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();
          new_body.push_back(MakeGroupBlock(
              block_stmt.size() == 1 ? block_stmt[0]
                                     : SeqStmt(std::move(block_stmt)),
              annotations));
755
756
        }
      }
757
    } else { // consumer case
758
759
      for (int i = 0; i < static_cast<int>(op->seq.size()); i++) {
        Array<Stmt> block_stmt = {};
760
761
        if (marker_.GetRole(op->seq[i]) == Role::kProducer)
          continue;
762
        for (int pattern_idx : map.acquire[i]) {
763
          PrimExpr acquire_barrier_id =
764
765
              stage_ + num_barriers_ + num_stages_ * pattern_idx;
          PrimExpr parity = map.is_loop_dependency(pattern_idx)
766
767
                                ? bitwise_xor(parity_, 1)
                                : parity_;
768
769
770
          block_stmt.push_back(makeParityWait(acquire_barrier_id, parity));
        }
        block_stmt.push_back(seq_transformed[i]);
771
772
773
774
775
        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;
776
            block_stmt.push_back(makeArriveBarrier(release_barrier_id));
777
778
779
780
            for (int s = 0; s < num_stages_; s++) {
              released_barrier_.insert(s + num_barriers_ +
                                       num_stages_ * pattern_idx);
            }
781
782
          }
        }
783
784
785
786
        new_body.push_back(MakeGroupBlock(block_stmt.size() == 1
                                              ? block_stmt[0]
                                              : SeqStmt(std::move(block_stmt)),
                                          annotations));
787
788
789
790
      }
      // Filter out the producer stmts
      int cur_id = 0;
      PipelineInfo new_pipeline_info;
791
792
      for (int i = 0; i < static_cast<int>(pipeline_info_.op_infos.size());
           i++) {
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
        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_;

812
    ICHECK(!new_body.empty());
813
814
815
    return new_body.size() == 1 ? new_body[0] : SeqStmt(std::move(new_body));
  }

816
  Stmt VisitStmt_(const ForNode *op) final {
817
818
    int num_stages = 1;
    auto num_stages_anno = op->annotations.Get("num_stages");
819
820
821
    if (num_stages_anno) {
      ICHECK(num_stages_anno->as<IntImmNode>());
      num_stages = static_cast<int>(num_stages_anno->as<IntImmNode>()->value);
822
823
      ICHECK(num_stages_ == 1) << "Nested pipeline not supported.";
    }
824
    loop_stack_.emplace_back(LoopInfo{op->loop_var, op->extent, op->min});
825
826
827
828

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

830
    auto group_anno = op->annotations.Get("tl_pipeline_group");
831
832
    if (group_anno) {
      group_info_array = Downcast<Array<Array<Integer>>>(group_anno.value());
833
834
    }
    auto order_anno = op->annotations.Get("tl_pipeline_order");
835
836
    if (order_anno) {
      order_info_array = Downcast<Array<Integer>>(order_anno.value());
837
838
    }
    auto stage_anno = op->annotations.Get("tl_pipeline_stage");
839
840
    if (stage_anno) {
      stage_info_array = Downcast<Array<Integer>>(stage_anno.value());
841
842
    }

843
844
    PipelineInfo pipeline_info(group_info_array, order_info_array,
                               stage_info_array);
845
846
    if (!pipeline_info.op_infos.empty()) {
      ICHECK(pipeline_info_.op_infos.empty())
847
          << "Nested pipeline not supported.";
848
849
850
851
852
853
854
855
856
    }

    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;
857
    PrimExpr linear_index = loop_stack_[0].loop_var - loop_stack_[0].min;
858
    for (size_t i = 1; i < loop_stack_.size(); ++i) {
859
860
      linear_index = linear_index * loop_stack_[i].extent +
                     (loop_stack_[i].loop_var - loop_stack_[i].min);
861
862
863
864
    }
    stage_ = FloorMod(linear_index, num_stages);
    parity_ = FloorMod(
        parity_before * op->extent + FloorDiv(linear_index, num_stages), 2);
865
866
867
    auto result = FilterByRole(op);

    Stmt grouped_for_node;
868
    if (result.as<ForNode>() && group_anno && !group_info_array.empty() &&
869
        !is_emitting_producer_) {
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
      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");
885
      if (is_emitting_producer_ || group_info_array.empty()) {
886
887
888
        for_node.CopyOnWrite()->annotations.erase("tl_pipeline_order");
        for_node.CopyOnWrite()->annotations.erase("tl_pipeline_stage");
      }
889
      if (is_emitting_producer_ || !group_anno || group_info_array.empty()) {
890
        loop_stack_.pop_back();
891
892
        return for_node;
      }
893
      loop_stack_.pop_back();
894
895
      return grouped_for_node;
    }
896
    loop_stack_.pop_back();
897
898
899
    return result;
  }

900
901
902
903
904
905
906
  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); }
  Stmt VisitStmt_(const BlockNode *op) final {
907
908
909
    ICHECK(0);
    return Stmt();
  }
910
  Stmt VisitStmt_(const BlockRealizeNode *op) final {
911
912
913
914
915
916
917
918
919
    ICHECK(0);
    return Stmt();
  }

  struct SyncPattern {
    int release_idx, acquire_idx;
  };

  struct SyncPatternMap {
920
921
922
    std::vector<std::vector<int>> acquire;
    std::vector<std::vector<int>> release;
    std::vector<std::vector<bool>> release_after;
923
    std::vector<SyncPattern> patterns;
924
925
926
927
928
929
930
931
932
933

    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;
934
935
936
    }
  };

937
  std::vector<SyncPattern>
938
  CreateBaseSyncPairs(const Array<Stmt> &seq_stmt,
939
                      const std::vector<bool> &is_producer) {
940
    const int n = seq_stmt.size();
941
    std::vector<std::set<const BufferNode *>> reads, writes;
942
943
944
    reads.reserve(n);
    writes.reserve(n);
    for (int i = 0; i < n; i++) {
945
946
      Block block(/*iter_vars=*/{}, /*reads=*/{}, /*writes=*/{},
                  /*name_hint=*/"",
947
948
                  /*body*/ seq_stmt[i]);
      auto access = GetBlockAccessRegion(block, buffer_data_to_buffer_);
949
      std::set<const BufferNode *> read_set, write_set;
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
      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());
        }
      }
966
967
968
969
      reads.push_back(std::move(read_set));
      writes.push_back(std::move(write_set));
    }

970
971
    auto intersect_fn = [](const std::set<const BufferNode *> &lhs,
                           const std::set<const BufferNode *> &rhs) {
972
      for (auto ptr : lhs)
973
974
        if (rhs.count(ptr))
          return true;
975
976
977
978
979
980
981
982
983
      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] &&
984
985
            (intersect_fn(writes[i], reads[j]) ||
             intersect_fn(reads[i], writes[j]))) {
986
987
988
989
990
991
992
993
994
995
996
997
998
999
          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] &&
1000
1001
              (intersect_fn(writes[i], reads[j]) ||
               intersect_fn(reads[i], writes[j]))) {
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
            sync_patterns.push_back({i, j});
            break;
          }
        }
      }
    }

    return sync_patterns;
  }

1012
1013
1014
  static std::vector<SyncPattern>
  RemoveUnusedSyncPatterns(const std::vector<SyncPattern> &sync_patterns,
                           const std::vector<bool> &is_producer) {
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
    /*
      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++)
1047
1048
      if (!removed[i])
        sync_pattern_cleaned.push_back(sync_patterns[i]);
1049
1050
1051
1052

    return sync_pattern_cleaned;
  }

1053
  SyncPatternMap ExtractSyncPattern(const Array<Stmt> &seq_stmt) {
1054
1055
1056
1057
1058
1059
1060
1061
    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);
1062
1063
    auto sync_patterns =
        RemoveUnusedSyncPatterns(sync_patterns_base, is_producer);
1064
1065

    // for (auto pattern : sync_patterns) {
1066
1067
    //   std::cout << pattern.release_idx << " " << pattern.acquire_idx <<
    //   std::endl;
1068
1069
1070
    // }

    SyncPatternMap map;
1071
    map.resize(num_stmts);
1072
    map.patterns = sync_patterns;
1073

1074
    for (size_t i = 0; i < sync_patterns.size(); i++) {
1075
1076
1077
1078
1079
1080
      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);
1081
1082
    }

1083
    std::vector<int> cur_consumer_barrier, cur_producer_barrier;
1084
1085
    for (int i = num_stmts - 1; i >= 0; i--) {
      if (is_producer[i]) {
1086
        if (map.release[i].empty()) {
1087
1088
1089
1090
          for (auto pattern_idx : cur_producer_barrier) {
            map.release[i].push_back(pattern_idx);
            map.release_after[i].push_back(false);
          }
1091
        } else {
1092
1093
1094
          for (auto pattern_idx : map.release[i]) {
            cur_producer_barrier.push_back(pattern_idx);
          }
1095
1096
        }
      } else {
1097
        if (map.release[i].empty()) {
1098
1099
1100
1101
          for (auto pattern_idx : cur_consumer_barrier) {
            map.release[i].push_back(pattern_idx);
            map.release_after[i].push_back(false);
          }
1102
        } else {
1103
1104
1105
          for (auto pattern_idx : map.release[i]) {
            cur_consumer_barrier.push_back(pattern_idx);
          }
1106
1107
1108
1109
1110
1111
1112
1113
1114
        }
      }
    }
    return map;
  }

  const bool is_emitting_producer_;
  Map<Var, Buffer> buffer_data_to_buffer_;
  std::unordered_set<int> released_barrier_;
1115
  const WarpSpecializedRoleMarker &marker_;
1116
1117
1118
1119
1120

  int num_barriers_ = 0;
  PrimExpr parity_ = 0;
  PrimExpr stage_ = 0;
  int num_stages_ = 1;
1121
  std::vector<LoopInfo> loop_stack_;
1122
  Var thread_var_;
1123
  bool mbarrier_only_ = false;
1124
1125
  PipelineInfo pipeline_info_;
  friend class WarpSpecializedRewriter;
1126
  bool has_simt_copy_ = false;
1127
1128
1129
};

class WarpSpecializedRewriter : public StmtExprMutator {
1130
public:
1131
1132
1133
1134
1135
1136
  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) {
1137
1138
1139
    // Check if function only uses threadIdx.x before proceeding
    if (!ThreadTagChecker::HasOnlyThreadIdxX(f)) {
      LOG(WARNING) << "WarpSpecialize will be disabled because the program "
1140
                      "uses thread tags other than threadIdx.x."
1141
1142
1143
1144
1145
1146
                   << "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;
    }

1147
1148
    auto T = WarpSpecializedRewriter(disable_warp_specialized,
                                     disable_shuffle_elect);
1149
    T.buffer_lca_ = DetectBufferAccessLCA(f);
1150
1151
    for (auto [buffer, _] : T.buffer_lca_)
      T.buffer_data_to_buffer_.Set(buffer->data, buffer);
1152
1153
1154
1155
    f.CopyOnWrite()->body = T(f->body);
    return f;
  }

1156
1157
private:
  Stmt VisitStmt_(const AttrStmtNode *op) final {
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
    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);
    }
  }

1175
1176
1177
1178
  // 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 {
1179
1180
1181
1182
1183
1184
1185
    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);
1186
      Stmt new_body =
1187
          ThreadIdxRewriter::Rewrite(for_node->body, thread_iv, thread_iv_, 0);
1188
1189
1190
1191
1192
      return new_body;
    }
    return for_node;
  }

1193
1194
1195
  Stmt VisitStmt_(const BlockRealizeNode *op) final {
    BlockRealize block_realize =
        Downcast<BlockRealize>(StmtExprMutator::VisitStmt_(op));
1196
1197
1198
1199
1200
1201
    if (!thread_iv_.defined()) {
      return block_realize;
    }

    Block block = block_realize->block;
    WarpSpecializedRoleMarker marker(buffer_data_to_buffer_);
1202
    marker.Prepare(block);
1203
1204
1205
1206
1207
1208
    marker(block);
    if (!marker.HasProducer()) {
      // Cannot detect any producer here, directly return.
      return block_realize;
    }

1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
    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(
1221
          DataType::Handle(), create_list_of_mbarrier(), barrier_num_threads));
1222
1223
1224
1225
      block.CopyOnWrite()->body = SeqStmt({init_barrier, code});
      block_realize.CopyOnWrite()->block = block;
      return block_realize;
    }
1226
    WSCodeEmitter producer(true, thread_iv_, buffer_data_to_buffer_, marker);
1227
    WSCodeEmitter consumer(false, thread_iv_, buffer_data_to_buffer_, marker,
1228
                           false);
1229
1230
1231
1232
1233
    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
1234
1235
    if (!marker.HasSimtCopy())
      producer_thread_extent = 128;
1236
1237

    updated_thread_extent_ = consumer_thread_extent + producer_thread_extent;
1238
1239
1240
1241
1242
1243
1244
1245

    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_);
1246
1247
1248
1249
1250
1251
1252
1253
    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++) {
1254
1255
1256
      PrimExpr arrive_thread_count =
          producer.released_barrier_.count(i)
              ? (producer.hasSimtCopy() ? producer_thread_extent : 1)
1257
              : consumer_thread_extent;
1258
1259
1260
      barrier_num_threads.push_back(arrive_thread_count);
    }

1261
    Stmt init_barrier = Evaluate(Call(
1262
        DataType::Handle(), create_list_of_mbarrier(), barrier_num_threads));
1263
1264
    Stmt body = IfThenElse(GE(thread_iv_->var, consumer_thread_extent),
                           producer_code, consumer_code);
1265
    // Add an attr here to handle the partial thread count in ThreadSync pass.
1266
1267
    Array<IntImm> ws_partition = {Downcast<IntImm>(producer_thread_extent),
                                  Downcast<IntImm>(consumer_thread_extent)};
1268
    body = AttrStmt(ws_partition, attr::kWarpSpecializationScope, 0, body);
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282

    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;
1283
  bool disable_warp_specialized_ = false;
1284
  bool disable_shuffle_elect_ = false;
1285
1286
};

1287
1288
class WarpSpecializedDetector : public IRVisitorWithAnalyzer {
public:
1289
  // return true means this aws will be disabled
1290
  static bool Detect(const Stmt &stmt, bool skip_thread_partition = false) {
1291
1292
    WarpSpecializedDetector detector;
    detector.VisitStmt(stmt);
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
    if (detector.has_warp_specialization_) {
      LOG(WARNING) << "Auto warp specialization will be disabled because warp "
                      "specialization is manually enabled";
      return true;
    }
    if (detector.has_tma_op_ && detector.has_mbarrier_op_) {
      LOG(WARNING) << "Auto warp specialization will be disabled because TMA "
                      "and mbarrier are both present";
      return true;
    }
    return false;
1304
1305
1306
1307
1308
  }

  WarpSpecializedDetector() {
    has_tma_op_ = false;
    has_mbarrier_op_ = false;
1309
    has_warp_specialization_ = false;
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
  }

private:
  void VisitStmt_(const EvaluateNode *op) final {
    if (const CallNode *call = op->value.as<CallNode>()) {
      if (call->op.same_as(create_list_of_mbarrier()) ||
          call->op.same_as(mbarrier_wait_parity()) ||
          call->op.same_as(builtin::ptx_arrive_barrier()) ||
          call->op.same_as(builtin::ptx_cp_async_barrier())) {
        has_mbarrier_op_ = true;
      }
    }
    IRVisitorWithAnalyzer::VisitStmt_(op);
  }

  void VisitExpr_(const CallNode *op) final {
    if (op->op.same_as(tma_load()) || op->op.same_as(tma_load_im2col()) ||
        op->op.same_as(set_max_nreg())) {
      has_tma_op_ = true;
    }
    IRVisitorWithAnalyzer::VisitExpr_(op);
  }

1333
  void VisitStmt_(const AttrStmtNode *op) final {
1334
1335
1336
1337
    if (op->attr_key == "warp_specialize" &&
        op->value.as<IntImmNode>()->value == 1) {
      has_warp_specialization_ = true;
    }
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
    if (op->attr_key == tir::attr::thread_extent) {
      IterVar iv = Downcast<IterVar>(op->node);
      if (iv->thread_tag == "threadIdx.x") {
        ICHECK(iv->dom->extent.as<IntImmNode>());
        thread_var_ = iv;
      }
    }
    IRVisitorWithAnalyzer::VisitStmt_(op);
  }

1348
  bool has_tma_op_{false};
1349
  IterVar thread_var_;
1350
  bool has_mbarrier_op_{false};
1351
  bool has_warp_specialization_{false};
1352
1353
};

1354
1355
1356
using namespace tir::transform;

tvm::transform::Pass WarpSpecialized() {
1357
  auto pass_func = [=](PrimFunc f, const IRModule &m, PassContext ctx) {
1358
1359
    bool disable_warp_specialized =
        ctx->GetConfig<Bool>(kDisableWarpSpecialized, Bool(false)).value();
1360
1361
    bool disable_shuffle_elect =
        ctx->GetConfig<Bool>(kDisableShuffleElect, Bool(false)).value();
1362
1363
1364
    bool warp_specialized = WarpSpecializedDetector::Detect(f->body);

    if (!warp_specialized) {
1365
1366
      return WarpSpecializedRewriter::Substitute(f, disable_warp_specialized,
                                                 disable_shuffle_elect);
1367
1368
    }
    return f;
1369
1370
1371
1372
  };
  return CreatePrimFuncPass(pass_func, 0, "tl.WarpSpecialized", {});
}

1373
1374
1375
1376
TVM_FFI_STATIC_INIT_BLOCK({
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def("tl.transform.WarpSpecialized", WarpSpecialized);
});
1377

1378
1379
} // namespace tl
} // namespace tvm