layout_inference.cc 49.2 KB
Newer Older
1
2
3
4
5
/*!
 * \file layout_inference.cc
 * \brief infer the fragment/shared memory layout
 */

6
#include <tvm/ffi/reflection/registry.h>
7
#include <tvm/tir/builtin.h>
8
#include <tvm/tir/index_map.h>
9
10
11
12
13
#include <tvm/tir/op.h>
#include <tvm/tir/stmt_functor.h>
#include <tvm/tir/transform.h>
#include <tvm/tir/utils.h>

14
#include <algorithm>
15
#include <deque>
16
#include <memory>
17
18
#include <queue>

19
#include "../layout/utils.h"
20
#include "../op/copy.h"
21
#include "../op/parallel.h"
22
#include "../op/region.h"
23
#include "../target/utils.h"
24

25
#include "arith/ir_mutator_with_analyzer.h"
26
#include "arith/ir_visitor_with_analyzer.h"
27
#include "common/loop_fusion_utils.h"
28
#include "common/loop_parallel_transform_utils.h"
29
#include "common/union_find.h"
30
#include "layout_reducer.h"
31
32
#include "loop_partition.h"
#include "loop_vectorize.h"
33
34
#include "runtime/thread_storage_scope.h"
#include "tir/transforms/ir_utils.h"
35
36
37
38

namespace tvm {
namespace tl {

39
40
41
using namespace tir;

/*!
42
 * \brief collect the mapping from the buffer var to it allocated buffer
43
 */
44
class ThreadBindingCollector : public StmtExprVisitor {
45
46
public:
  void VisitStmt_(const AttrStmtNode *op) final {
47
48
49
50
    if (op->attr_key == tir::attr::thread_extent) {
      IterVar iv = Downcast<IterVar>(op->node);
      thread_binding_[iv->var.get()] = iv;
    }
51
52
53
    StmtExprVisitor::VisitStmt_(op);
  }

54
55
  // The thread binding map
  std::unordered_map<const VarNode *, IterVar> thread_binding_;
56
57
};

58
59
using namespace tir;
using arith::IRMutatorWithAnalyzer;
60
using arith::IRVisitorWithAnalyzer;
61
62
63
64
65
66
67

struct LayoutInferenceResult {
  Map<Buffer, Layout> layout_map;
  Map<For, Fragment> for_map;
  Map<For, PrimExpr> predicate_map;
};

68
class BufferUseDefCollector : public IRVisitorWithAnalyzer {
69
public:
70
71
  BufferUseDefCollector(bool skip_thread_partition)
      : skip_thread_partition_(skip_thread_partition) {}
72

73
74
  using arith::IRVisitorWithAnalyzer::IRVisitorWithAnalyzer;

75
76
  void RunInferStep(int cur_infer_id, InferLevel level, bool update_queue,
                    LayoutMap &layout_map, const LayoutMap &strict_layout_map,
77
                    std::deque<int> &q, std::vector<bool> &in_queue) {
78
79
80
81
82
83
84
85
86
87
88
89
90
    auto num_infer = infer_list_.size();

    // Range check for cur_infer_id
    ICHECK_GE(cur_infer_id, 0) << "cur_infer_id is negative, which is invalid.";
    ICHECK_LT(cur_infer_id, num_infer)
        << "cur_infer_id " << cur_infer_id << " is out of range, must be < "
        << num_infer << ".";

    // Make sure we can safely access infer_list_[cur_infer_id] and
    // thread_var_vec_[cur_infer_id]
    auto &next = infer_list_[cur_infer_id];
    auto iter_var = thread_var_vec_[cur_infer_id];
    auto thread_bounds = thread_bounds_vec_[cur_infer_id];
91
    arith::Analyzer *cur_analyzer = analyzer_vec_[cur_infer_id].get();
92
    auto buffer_oob = buffer_oob_vec_[cur_infer_id];
93
    // Double-check that 'next' is valid
94
95
    ICHECK(next.defined()) << "infer_list_[" << cur_infer_id
                           << "] is null inside run_infer_step.";
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112

    // Check iter_var->dom and dom->extent
    ICHECK(iter_var.defined())
        << "thread_var_vec_[" << cur_infer_id << "] is not defined.";
    ICHECK(iter_var->dom.defined())
        << "iter_var->dom is not defined for infer_list_[" << cur_infer_id
        << "].";
    ICHECK(iter_var->dom->extent.defined())
        << "iter_var->dom->extent is not defined for infer_list_["
        << cur_infer_id << "].";

    const int64_t *extent_ptr = as_const_int(iter_var->dom->extent);
    ICHECK(extent_ptr != nullptr)
        << "iter_var->dom->extent is not a constant integer, which is "
           "required for layout inference.";

    // Run InferLayout
113
114
115
116
117
118
119
120
    auto updates = next->InferLayout(LayoutInferArgs{target_,
                                                     thread_bounds,
                                                     layout_map,
                                                     cur_analyzer,
                                                     buffer_oob,
                                                     {},
                                                     let_var_to_expr_},
                                     level);
121

122
123
124
125
126
127
    // Process the returned updates
    for (const auto &[buffer, layout] : updates) {
      // Basic validity checks
      ICHECK(buffer.defined()) << "InferLayout returned an undefined buffer.";
      ICHECK(layout.defined()) << "InferLayout returned an undefined layout.";

128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
      // Helper: propagate inferred layout to alias buffers (same data Var)
      auto propagate_alias = [&](const Buffer &src_buffer,
                                 const Layout &src_layout) {
        if (!buffer_data_to_buffers_.count(src_buffer->data))
          return;
        const auto &siblings = buffer_data_to_buffers_[src_buffer->data];
        for (const auto &sib : siblings) {
          if (sib.same_as(src_buffer))
            continue;
          bool shapes_equal =
              src_layout->InputShape().size() == sib->shape.size();
          if (shapes_equal) {
            for (size_t i = 0; i < src_layout->InputShape().size(); ++i) {
              if (!analyzer_.CanProveEqual(src_layout->InputShape()[i],
                                           sib->shape[i])) {
                shapes_equal = false;
                break;
              }
            }
          }
          Layout target_layout =
149
150
151
152
153
              shapes_equal
                  ? src_layout
                  : src_layout->Reshape(sib->shape, &analyzer_,
                                        Integer(src_buffer->dtype.bytes()),
                                        Integer(sib->dtype.bytes()));
154
155
156
157
158
159
160
161
162
163
          if (layout_map.count(sib)) {
            ICHECK(target_layout->IsEqual(layout_map[sib].get()))
                << "Get different layout for alias buffer " << sib
                << " (data-shared with " << src_buffer
                << ")\n current: " << target_layout->DebugOutput()
                << "\n previous: " << layout_map[sib]->DebugOutput();
          } else {
            layout_map.Set(sib, target_layout);
            if (update_queue && use_list_.count(sib)) {
              for (int idx : use_list_[sib]) {
164
                EnqueueWithPriority(idx, q, in_queue, cur_infer_id, layout_map);
165
166
167
168
169
170
              }
            }
          }
        }
      };

171
172
173
174
175
176
      if (layout_map.count(buffer)) {
        // If new layout contains the old one, update map
        if (buffer.scope() == "local.fragment" &&
            level != InferLevel::kStrict && !strict_layout_map.count(buffer)) {
          // Actually this test has been done in ParallelOp::InferLayout
          // already. Just do it again to avoid missing implementations in other
177
          // `TileOperator`s.
178
179
180
181
182

          auto dst_layout_opt = layout.as<Fragment>();
          ICHECK(dst_layout_opt.has_value())
              << "Failed to cast layout to Fragment for buffer " << buffer
              << ", layout type is " << layout->GetTypeKey();
183
          const auto &dst_layout = dst_layout_opt.value();
184
185
186
187
188
          auto src_layout_opt = layout_map[buffer].as<Fragment>();
          ICHECK(src_layout_opt.has_value())
              << "Failed to cast layout_map[buffer] to Fragment for buffer "
              << buffer << ", layout type is "
              << layout_map[buffer]->GetTypeKey();
189
          const auto &src_layout = src_layout_opt.value();
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
          ICHECK(dst_layout->InputDim() == src_layout->InputDim());
          Array<PrimExpr> indices;
          indices.reserve(dst_layout->InputDim());
          arith::Analyzer inner_analyzer;
          for (int i = 0; i < dst_layout->InputDim(); ++i) {
            auto x = InputPlaceholder(i);
            indices.push_back(x);
            // should be literal - literal = 0, any analyzer will work
            ICHECK(is_zero(inner_analyzer.Simplify(
                dst_layout->InputShape()[i] - src_layout->InputShape()[i])));
            inner_analyzer.Bind(x, Range(0, dst_layout->InputShape()[i]));
          }
          if (ProveFragmentContains(src_layout, dst_layout, indices, indices,
                                    inner_analyzer)) {
            layout_map.Set(buffer, layout);
205
206
            // Propagate to alias buffers as well
            propagate_alias(buffer, layout);
207
208
209
210
            continue;
          }
        }
        // If already in map, ensure they are structurally equal
211
        ICHECK(layout->IsEqual(layout_map[buffer].get()))
212
213
214
            << "Get different layout for " << buffer
            << "\n current layout: " << layout->DebugOutput()
            << "\n previous layout: " << layout_map[buffer]->DebugOutput();
215
216
        // Ensure aliases are consistent too
        propagate_alias(buffer, layout);
217
218
219
      } else {
        // Otherwise, update map
        layout_map.Set(buffer, layout);
220
221
        // Propagate to alias buffers (may enqueue their users)
        propagate_alias(buffer, layout);
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
        if (!update_queue)
          continue;

        // Check if buffer exists in use_list_
        if (!use_list_.count(buffer)) {
          LOG(WARNING) << "Layout inference failed for buffer " << buffer
                       << ". "
                       << "The buffer cannot be inferred with current layout "
                          "inference rules.";
          continue;
        }

        // Push back into BFS queue
        for (int idx : use_list_[buffer]) {
          ICHECK_GE(idx, 0)
              << "Index in use_list_ for buffer " << buffer << " is negative.";
          ICHECK_LT(idx, num_infer)
              << "Index in use_list_ for buffer " << buffer
              << " out of range: " << idx << " >= " << num_infer << ".";

242
          EnqueueWithPriority(idx, q, in_queue, cur_infer_id, layout_map);
243
244
245
246
247
248
        }
      }
    }
  };

  void FinishInferQueue(InferLevel level, LayoutMap &layout_map,
249
                        const LayoutMap &strict_layout_map, std::deque<int> &q,
250
251
                        std::vector<bool> &in_queue) {
    auto num_infer = infer_list_.size();
252

253
254
    while (!q.empty()) {
      int cur_infer_id = q.front();
255
      q.pop_front();
256
257
258
259
260
261
262
263
264
265
      // Range check again, just to be safe
      ICHECK_GE(cur_infer_id, 0);
      ICHECK_LT(cur_infer_id, num_infer);

      in_queue[cur_infer_id] = false;
      RunInferStep(cur_infer_id, level, true, layout_map, strict_layout_map, q,
                   in_queue);
    }
  };

266
  LayoutInferenceResult Run() {
267
268
269
270
271
    // Basic consistency check: infer_list_ and thread_var_vec_ should have the
    // same size
    ICHECK_EQ(infer_list_.size(), thread_var_vec_.size())
        << "Size mismatch: infer_list_ and thread_var_vec_ must match in "
           "length.";
272
273
274
    ICHECK_EQ(thread_bounds_vec_.size(), infer_list_.size())
        << "Size mismatch: thread_bounds_vec_ and infer_list_ must match in "
           "length.";
275
276
277
    ICHECK_EQ(analyzer_vec_.size(), infer_list_.size())
        << "Size mismatch: analyzer_vec_ and infer_list_ must match in "
           "length.";
278
279
280
    ICHECK_EQ(buffer_oob_vec_.size(), infer_list_.size())
        << "Size mismatch: buffer_oob_vec_ and infer_list_ must match in "
           "length.";
281

282
283
284
285
286
    DLOG(INFO) << "[InferLayout] all participating operators:" << '\n';
    for (int i = 0; i < infer_list_stmt_.size(); ++i) {
      DLOG(INFO) << "    op " << i << ":" << infer_list_stmt_[i] << '\n';
    }

287
288
289
290
    // If needed, you can also check that annotated_layout_map_ is not empty, or
    // anything else relevant to your setup.

    // Copy the annotated layout map to local variable
291
    Map<Buffer, Layout> layout_map = annotated_layout_map_;
292
    Map<Buffer, Layout> strict_layout_map;
293
294
    int num_infer = infer_list_.size();

295
    // Prepare BFS queue for iterative inference
296
    std::deque<int> q;
297
    std::vector<bool> in_queue(num_infer, true);
298
299
    for (int i = 0; i < num_infer; i++) {
      // Check that each infer_list_ entry is valid
300
      ICHECK(infer_list_[i].defined())
301
302
303
304
305
306
307
          << "infer_list_[" << i
          << "] is null. The inference object is not allocated properly.";

      // Check that each thread_var_vec_ entry is defined
      if (!thread_var_vec_[i].defined() && skip_thread_partition_) {
        thread_var_vec_[i] = thread_var_;
      }
308
      q.push_back(i);
309
    }
310

311
    // step 1: infer strict layout
312
    for (int i = 0; i < num_infer; i++) {
313
314
      RunInferStep(i, InferLevel::kStrict, false, layout_map, strict_layout_map,
                   q, in_queue);
315
316
    }

317
318
319
320
    for (const auto &[buffer, layout] : layout_map) {
      strict_layout_map.Set(buffer, layout);
    }

321
    // step 2: infer common layout with BFS
322
323
    FinishInferQueue(InferLevel::kCommon, layout_map, strict_layout_map, q,
                     in_queue);
324

325
    // step 3: relax constraints to free and re-run
326
327
    InferInFreeMode(layout_map, strict_layout_map);

328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
    // step 4: finalize alias layouts by Var
    // For each storage var, if any buffer in the group has a layout,
    // propagate (reshape if needed) to the rest to ensure completeness.
    for (const auto &[var, buffers] : buffer_data_to_buffers_) {
      // Find a representative with existing layout
      Optional<Buffer> rep;
      Optional<Layout> rep_layout;
      for (const auto &buf : buffers) {
        if (layout_map.count(buf)) {
          rep = buf;
          rep_layout = layout_map[buf];
          break;
        }
      }
      if (!rep_layout.defined())
        continue;
      for (const auto &buf : buffers) {
        if (!layout_map.count(buf)) {
          bool shapes_equal =
              rep_layout.value()->InputShape().size() == buf->shape.size();
          if (shapes_equal) {
            for (size_t i = 0; i < rep_layout.value()->InputShape().size();
                 ++i) {
              if (!analyzer_.CanProveEqual(rep_layout.value()->InputShape()[i],
                                           buf->shape[i])) {
                shapes_equal = false;
                break;
              }
            }
          }

359
360
361
362
363
364
          Layout reshaped = shapes_equal
                                ? rep_layout.value()
                                : rep_layout.value()->Reshape(
                                      buf->shape, &analyzer_,
                                      Integer(rep.value()->dtype.bytes()),
                                      Integer(buf->dtype.bytes()));
365
366
367
368
369
          layout_map.Set(buf, reshaped);
        }
      }
    }

370
    // Check that all local.fragment buffers have inferred layouts
371
    for (const auto &[buffer, _] : use_list_) {
372
373
374
375
376
      if (buffer.scope() == "local.fragment") {
        ICHECK_NE(layout_map.count(buffer), 0)
            << "The layout for fragment " << buffer
            << " can not be inferred correctly.";
      }
377
378
    }

379
    // Collect layout info for For nodes
380
381
    Map<For, Fragment> for_map;
    Map<For, PrimExpr> predicate_map;
382
383
384
    ICHECK(infer_list_.size() == thread_var_vec_.size())
        << "infer_list_ and thread_var_vec_ size mismatch";
    for (int i = 0; i < infer_list_.size(); i++) {
385
      TileOperator base_infer = std::move(infer_list_[i]);
386
387
      auto thread_var = thread_var_vec_[i];

388
      // Check if base_infer is valid
389
390
391
      ICHECK(base_infer.defined()) << "Null pointer encountered in "
                                      "infer_list_ while collecting for_map.";
      if (auto for_infer = base_infer.as<ParallelOpNode>()) {
392
        // Check that the loop layout is defined
393
        ICHECK(for_infer->GetLoopLayout().defined())
394
            << "The Layout for Parallel for cannot be inferred correctly:\n"
395
396
            << for_infer->GetRoot();
        for_map.Set(for_infer->GetRoot(), for_infer->GetLoopLayout());
397
        // thread_var_ should be defined if we rely on it
398
399
        ICHECK(thread_var.defined())
            << "thread_var is not defined. Cannot retrieve predicate.";
400

401
        if (auto predicate = for_infer->GetPredicate(thread_var->var)) {
402
          predicate_map.Set(for_infer->GetRoot(), predicate.value());
403
        }
404
405
406
407
408
409
      }
    }

    return {layout_map, for_map, predicate_map};
  }

410
411
  void Collect(const PrimFunc &f) {
    for (const auto &[_, buffer] : f->buffer_map) {
412
413
414
415
416
417
418
      if (buffer_data_to_buffers_.count(buffer->data)) {
        auto buffers = buffer_data_to_buffers_[buffer->data];
        buffers.push_back(buffer);
        buffer_data_to_buffers_.Set(buffer->data, buffers);
      } else {
        buffer_data_to_buffers_.Set(buffer->data, {buffer});
      }
419
420
    }
    auto target = f->GetAttr<Target>(tvm::attr::kTarget);
421
422
    ICHECK(target.defined())
        << "Layout_Inference: Require the target attribute";
423
424
425
426
    target_ = target.value();
    this->operator()(f->body);
  }

427
private:
428
429
430
431
432
433
434
435
436
437
438
439
  Map<Var, Buffer> GetBufferMap() const {
    Map<Var, Buffer> buffer_map;
    for (const auto &[var, buffers] : buffer_data_to_buffers_) {
      // Use the first buffer for each var
      // TODO(lei): phaseout buffer_map in future.
      if (!buffers.empty()) {
        buffer_map.Set(var, buffers[0]);
      }
    }
    return buffer_map;
  }

440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
  // Return true if all buffers that this op (idx) touches already have
  // inferred layouts in layout_map. Used to prioritize enqueue order.
  bool ShouldPrioritize(int idx, const LayoutMap &layout_map) const {
    auto it = op_touched_buffers_.find(idx);
    if (it == op_touched_buffers_.end() || it->second.empty())
      return false;
    for (const auto &buf : it->second) {
      if (!layout_map.count(buf))
        return false;
    }
    return true;
  }

  // Enqueue idx to q with priority if all its buffers already
  // have layouts. Also guards against duplicates and self-enqueue.
  void EnqueueWithPriority(int idx, std::deque<int> &q,
                           std::vector<bool> &in_queue, int cur_infer_id,
                           const LayoutMap &layout_map) const {
    if (idx == cur_infer_id)
      return;
    if (idx < 0 || idx >= static_cast<int>(in_queue.size()))
      return;
    if (in_queue[idx])
      return;
    in_queue[idx] = true;
    if (ShouldPrioritize(idx, layout_map)) {
      q.push_front(idx);
    } else {
      q.push_back(idx);
    }
  }

472
  void VisitExpr_(const CallNode *op) final {
473
    IRVisitorWithAnalyzer::VisitExpr_(op);
474
    // Do not analysis the call node to the global function.
475
476
    if (op->op.as<GlobalVarNode>())
      return;
477

478
    auto p = ParseOperator(tvm::ffi::GetRef<Call>(op));
479
    if (p.defined()) {
480
      for (const auto &arg : op->args) {
481
482
        if (auto buffer = getBufferFromAccessPtr(arg)) {
          addToUseList(buffer.value());
483
484
        } else if (auto buffer = getBufferFromRegion(arg)) {
          addToUseList(buffer.value());
485
        }
486
487
488
489
        // Check if the argument uses any LetStmt variables that reference
        // fragment buffers. If so, add those buffers to the use list.
        // This handles cases like: a = block_mask_f[i]; T.copy(A[a, 0], ...)
        CollectFragmentBuffersFromExpr(arg);
490
      }
491
      // Compute thread_var_ and thread_bounds_
492
      thread_var_vec_.push_back(thread_var_);
493
494
495
496
      if (analyzer_.const_int_bound.IsBound(thread_var_->var)) {
        auto const_int_bound = analyzer_.const_int_bound(thread_var_);
        auto min_value = const_int_bound->min_value;
        auto max_value = const_int_bound->max_value;
497
        auto extent = max_value - min_value + 1;
498
499
        auto dtype = thread_var_->var.dtype();
        thread_bounds_vec_.push_back(Range::FromMinExtent(
500
            IntImm(dtype, min_value), IntImm(dtype, extent)));
501
502
503
      } else {
        thread_bounds_vec_.push_back(Range::FromMinExtent(0, 1));
      }
504
      analyzer_vec_.push_back(analyzer_.Clone());
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535

      // Compute buffer oob for each buffer in the op
      if (const auto *copy = p.as<CopyNode>()) {
        auto src_tensor = copy->src;
        auto dst_tensor = copy->dst;
        auto src_range = copy->src_range;
        auto dst_range = copy->dst_range;
        bool src_oob = false;
        bool dst_oob = false;
        for (size_t i = 0; i < src_range.size(); i++) {
          if (!analyzer_.CanProve(src_range[i]->min + src_range[i]->extent <=
                                      src_tensor->shape[i],
                                  arith::ProofStrength::kSymbolicBound)) {
            src_oob = true;
            break;
          }
        }
        for (size_t i = 0; i < dst_range.size(); i++) {
          if (!analyzer_.CanProve(dst_range[i]->min + dst_range[i]->extent <=
                                      dst_tensor->shape[i],
                                  arith::ProofStrength::kSymbolicBound)) {
            dst_oob = true;
            break;
          }
        }
        buffer_oob_vec_.push_back(src_oob || dst_oob);
      } else {
        buffer_oob_vec_.push_back(false);
      }

      // Add the tile operator to infer_list_
536
      infer_list_stmt_.push_back(tvm::ffi::GetRef<ObjectRef>(op));
537
      infer_list_.push_back(std::move(p));
538
539
540
    }
  }

541
  Optional<Buffer> getBufferFromAccessPtr(const PrimExpr &expr) {
542
543
544
    if (auto bl = expr.as<BufferLoadNode>()) {
      return bl->buffer;
    }
545
    auto call = expr.as<CallNode>();
546
547
548
549
    if (!call) {
      return std::nullopt;
    }
    if (call->op.same_as(builtin::tvm_access_ptr())) {
550
551
      auto var_opt = call->args[1].as<Var>();
      if (!var_opt.has_value()) {
552
553
        LOG(WARNING) << "[getBufferFromAccessPtr] args[1] is not a Var, type: "
                     << call->args[1]->GetTypeKey();
554
555
        return std::nullopt;
      }
556
      const auto &var = var_opt.value();
557
558
559
560
561
562
563
      if (buffer_data_to_buffers_.count(var)) {
        const auto &buffers = buffer_data_to_buffers_[var];
        if (!buffers.empty()) {
          return buffers[0]; // Return the first buffer
        }
      }
      return std::nullopt;
564
565
566
567
568
569
570
571
572
573
574
575
    }
    return std::nullopt;
  }

  Optional<Buffer> getBufferFromRegion(const PrimExpr &expr) {
    if (auto call = expr.as<CallNode>()) {
      if (call->op.same_as(RegionOp::Get())) {
        if (auto bl = call->args[0].as<BufferLoadNode>()) {
          return bl->buffer;
        }
        return std::nullopt;
      }
576
    }
577
    return std::nullopt;
578
579
  }

580
  void addToUseList(const Buffer &buffer) {
581
582
583
584
    // buffer scope must be local.fragment
    if (buffer.scope() != "local.fragment") {
      return;
    }
585
586
587
588
589
    int infer_idx = infer_list_.size();
    if (use_list_.find(buffer) == use_list_.end()) {
      use_list_[buffer] = {};
    }
    use_list_[buffer].push_back(infer_idx);
590
591
592
593
594
595
596
597
598
599
600
601
602

    // Track which buffers this op (infer_idx) touches for prioritization.
    // Avoid duplicates.
    auto &vec = op_touched_buffers_[infer_idx];
    bool exists = false;
    for (const auto &b : vec) {
      if (b.same_as(buffer)) {
        exists = true;
        break;
      }
    }
    if (!exists)
      vec.push_back(buffer);
603
604
  }

605
  void VisitStmt_(const ForNode *op) final {
606
    if (op->kind == ForKind::kParallel) {
607
      auto infer = ParallelOp(tvm::ffi::GetRef<For>(op));
608
      for (const auto &[buffer, _] : infer->GetIndiceMap()) {
609
610
        addToUseList(buffer);
      }
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
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
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675

      PostOrderVisit(op->body, [this](const ObjectRef &node) {
        if (auto *buffer_load = node.as<BufferLoadNode>()) {
          if (buffer_load->buffer.defined() &&
              buffer_load->buffer->data.defined()) {
            if (buffer_data_to_buffers_.count(buffer_load->buffer->data)) {
              // Check if this buffer is already in the list
              auto buffers = buffer_data_to_buffers_[buffer_load->buffer->data];
              bool found = false;
              for (const auto &buf : buffers) {
                if (buf.same_as(buffer_load->buffer)) {
                  found = true;
                  break;
                }
              }
              if (!found) {
                buffers.push_back(buffer_load->buffer);
                buffer_data_to_buffers_.Set(buffer_load->buffer->data, buffers);
                DLOG(INFO) << "[LayoutInference] BufferStore: added buffer "
                           << buffer_load->buffer
                           << " buffer.get() = " << buffer_load->buffer.get()
                           << " data = " << buffer_load->buffer->data.get();
              }
            } else {
              buffer_data_to_buffers_.Set(buffer_load->buffer->data,
                                          {buffer_load->buffer});
              DLOG(INFO) << "[LayoutInference] BufferStore: new buffer "
                         << buffer_load->buffer
                         << " buffer.get() = " << buffer_load->buffer.get()
                         << " data = " << buffer_load->buffer->data.get();
            }
          }
        } else if (auto *buffer_store = node.as<BufferStoreNode>()) {
          if (buffer_store->buffer.defined() &&
              buffer_store->buffer->data.defined()) {
            if (buffer_data_to_buffers_.count(buffer_store->buffer->data)) {
              auto buffers =
                  buffer_data_to_buffers_[buffer_store->buffer->data];
              bool found = false;
              for (const auto &buf : buffers) {
                if (buf.same_as(buffer_store->buffer)) {
                  found = true;
                  break;
                }
              }
              if (!found) {
                buffers.push_back(buffer_store->buffer);
                buffer_data_to_buffers_.Set(buffer_store->buffer->data,
                                            buffers);
                DLOG(INFO) << "[LayoutInference] BufferStore: added buffer "
                           << buffer_store->buffer
                           << " buffer.get() = " << buffer_store->buffer.get()
                           << " data = " << buffer_store->buffer->data.get();
              }
            } else {
              buffer_data_to_buffers_.Set(buffer_store->buffer->data,
                                          {buffer_store->buffer});
              DLOG(INFO) << "[LayoutInference] BufferStore: new buffer "
                         << buffer_store->buffer
                         << " buffer.get() = " << buffer_store->buffer.get()
                         << " data = " << buffer_store->buffer->data.get();
            }
          }
        }
      });
676
      infer_list_stmt_.push_back(tvm::ffi::GetRef<ObjectRef>(op));
677
678
      infer_list_.push_back(std::move(infer));
      thread_var_vec_.push_back(thread_var_);
679
680
681
682
      if (thread_var_.defined() &&
          analyzer_.const_int_bound.IsBound(thread_var_->var)) {
        auto const_int_bound = analyzer_.const_int_bound(thread_var_);
        auto dtype = thread_var_->var.dtype();
683
684
        auto extent =
            const_int_bound->max_value - const_int_bound->min_value + 1;
685
        thread_bounds_vec_.push_back(Range::FromMinExtent(
686
            IntImm(dtype, const_int_bound->min_value), IntImm(dtype, extent)));
687
688
689
      } else {
        thread_bounds_vec_.push_back(Range::FromMinExtent(0, 1));
      }
690
      analyzer_vec_.push_back(analyzer_.Clone());
691
      buffer_oob_vec_.push_back(false);
692
    } else {
693
      IRVisitorWithAnalyzer::VisitStmt(op->body);
694
695
696
    }
  }

697
  void VisitStmt_(const BlockNode *op) final {
698
    for (auto buffer : op->alloc_buffers) {
699
700
701
702
703
704
705
      if (buffer_data_to_buffers_.count(buffer->data)) {
        auto buffers = buffer_data_to_buffers_[buffer->data];
        buffers.push_back(buffer);
        buffer_data_to_buffers_.Set(buffer->data, buffers);
      } else {
        buffer_data_to_buffers_.Set(buffer->data, {buffer});
      }
706
    }
707
708
709
710
711
712

    // First, visit the block body to collect all buffers from
    // BufferLoad/BufferStore
    IRVisitorWithAnalyzer::VisitStmt_(op);

    // After visiting, apply layouts to all collected buffers
713
    if (op->annotations.count(attr::kLayoutMap)) {
714
      // Check if the layout map is Map<Var, Layout>
715
716
717
      auto map =
          op->annotations.Get(attr::kLayoutMap)->as<Map<Var, Layout>>().value();
      for (const auto &[var, layout] : map) {
718
        ICHECK(buffer_data_to_buffers_.count(var))
719
            << "buffer " << var << " is not found in the block";
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
        const auto &buffers = buffer_data_to_buffers_[var];
        ICHECK(!buffers.empty()) << "buffer list for " << var << " is empty";
        // Apply layout to all buffers associated with this var
        for (const auto &buffer : buffers) {

          // Reshape the layout to match the buffer's shape
          // Check if shapes are structurally equal
          bool shapes_equal =
              layout->InputShape().size() == buffer->shape.size();
          if (shapes_equal) {
            for (size_t i = 0; i < layout->InputShape().size(); ++i) {
              if (!analyzer_.CanProveEqual(layout->InputShape()[i],
                                           buffer->shape[i])) {
                shapes_equal = false;
                break;
              }
            }
          }

          if (shapes_equal) {
            annotated_layout_map_.Set(buffer, layout);
          } else {
742
743
744
745
746
            // Use the first buffer sharing this var as the base for dtype ratio
            int base_bytes = buffers[0]->dtype.bytes();
            auto reshaped_layout =
                layout->Reshape(buffer->shape, &analyzer_, Integer(base_bytes),
                                Integer(buffer->dtype.bytes()));
747
748
749
            annotated_layout_map_.Set(buffer, reshaped_layout);
          }
        }
750
751
752
753
      }
    }
  }

754
  void VisitStmt_(const AttrStmtNode *op) final {
755
756
757
758
759
760
761
    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;
      }
    }
762
    IRVisitorWithAnalyzer::VisitStmt_(op);
763
764
  }

765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
  void VisitStmt_(const LetStmtNode *op) final {
    // Record Let variable to its bound expression.
    // This enables tracking fragment buffer accesses through let bindings.
    let_var_to_expr_.Set(op->var, op->value);
    IRVisitorWithAnalyzer::VisitStmt_(op);
  }

  // Helper: recursively collect fragment buffers from an expression,
  // following let bindings chain.
  void CollectFragmentBuffersFromExpr(const PrimExpr &expr) {
    PostOrderVisit(expr, [this](const ObjectRef &node) {
      if (auto bl = node.as<BufferLoadNode>()) {
        if (bl->buffer.defined() && bl->buffer.scope() == "local.fragment") {
          addToUseList(bl->buffer);
        }
      } else if (auto var_node = node.as<VarNode>()) {
        auto var = tvm::ffi::GetRef<Var>(var_node);
        if (let_var_to_expr_.count(var)) {
          CollectFragmentBuffersFromExpr(let_var_to_expr_[var]);
        }
      }
    });
  }

789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
  void VisitExpr_(const BufferLoadNode *op) final {
    // Collect buffer from BufferLoad
    if (op->buffer.defined() && op->buffer->data.defined()) {
      if (buffer_data_to_buffers_.count(op->buffer->data)) {
        // Check if this buffer is already in the list
        auto buffers = buffer_data_to_buffers_[op->buffer->data];
        bool found = false;
        for (const auto &buf : buffers) {
          if (buf.same_as(op->buffer)) {
            found = true;
            break;
          }
        }
        if (!found) {
          buffers.push_back(op->buffer);
          buffer_data_to_buffers_.Set(op->buffer->data, buffers);
          DLOG(INFO) << "[LayoutInference] BufferLoad: added buffer "
                     << op->buffer << " buffer.get() = " << op->buffer.get()
                     << " data = " << op->buffer->data.get();
        }
      } else {
        buffer_data_to_buffers_.Set(op->buffer->data, {op->buffer});
        DLOG(INFO) << "[LayoutInference] BufferLoad: new buffer " << op->buffer
                   << " buffer.get() = " << op->buffer.get()
                   << " data = " << op->buffer->data.get();
      }
    }
    IRVisitorWithAnalyzer::VisitExpr_(op);
  }

  void VisitStmt_(const BufferStoreNode *op) final {
    // Collect buffer from BufferStore
    if (op->buffer.defined() && op->buffer->data.defined()) {
      if (buffer_data_to_buffers_.count(op->buffer->data)) {
        // Check if this buffer is already in the list
        auto buffers = buffer_data_to_buffers_[op->buffer->data];
        bool found = false;
        for (const auto &buf : buffers) {
          if (buf.same_as(op->buffer)) {
            found = true;
            break;
          }
        }
        if (!found) {
          buffers.push_back(op->buffer);
          buffer_data_to_buffers_.Set(op->buffer->data, buffers);
          DLOG(INFO) << "[LayoutInference] BufferStore: added buffer "
                     << op->buffer << " buffer.get() = " << op->buffer.get()
                     << " data = " << op->buffer->data.get();
        }
      } else {
        buffer_data_to_buffers_.Set(op->buffer->data, {op->buffer});
        DLOG(INFO) << "[LayoutInference] BufferStore: new buffer " << op->buffer
                   << " buffer.get() = " << op->buffer.get()
                   << " data = " << op->buffer->data.get();
      }
    }
    IRVisitorWithAnalyzer::VisitStmt_(op);
  }

  Map<Var, Array<Buffer>> buffer_data_to_buffers_;
850
851
  // Map from LetStmt variable to its bound expression
  Map<Var, PrimExpr> let_var_to_expr_;
852
  std::vector<ObjectRef> infer_list_stmt_;
853
  std::vector<TileOperator> infer_list_;
854
855
  std::unordered_map<Buffer, std::vector<int>, ObjectPtrHash, ObjectPtrEqual>
      use_list_;
856
857
  // Per-op list of buffers it touches (fragment scope), used for prioritization
  std::unordered_map<int, std::vector<Buffer>> op_touched_buffers_;
858
859
860
861
  // This is a workaround for cpu backend,
  // we need to define a thread_var for the serial loop.
  IterVar thread_var_ = IterVar(Range::FromMinExtent(0, 1), Var("v_thread"),
                                IterVarType::kDataPar);
862
  std::vector<IterVar> thread_var_vec_;
863
  std::vector<Range> thread_bounds_vec_;
864
  std::vector<std::unique_ptr<arith::Analyzer>> analyzer_vec_;
865
  std::vector<bool> buffer_oob_vec_;
866
867
  Target target_;
  LayoutMap annotated_layout_map_;
868
  bool skip_thread_partition_{false};
869

870
871
  std::vector<TileOperator> BackupInferList() {
    std::vector<TileOperator> back_infer_list;
872
873
874
875
876
877
878
879
880
    back_infer_list.reserve(infer_list_.size());
    for (auto &&p : infer_list_) {
      back_infer_list.push_back(p->Clone());
    }
    return back_infer_list;
  }

  void InferInFreeMode(LayoutMap &layout_map,
                       const LayoutMap &strict_layout_map) {
881
882
883
884
885
886
887

    DLOG(INFO) << "Enforced layout maps:" << '\n';
    for (auto &&[k, v] : layout_map) {
      DLOG(INFO) << "    " << k << ": " << v->DebugOutput() << '\n';
    }
    DLOG(INFO) << '\n';

888
889
890
891
892
893
894
895
896
    // Group operators into connected components
    UnionFind<int> uf;
    for (int i = 0; i < infer_list_.size(); i++) {
      uf.MakeSet(i);
    }
    for (const auto &[buffer, infer_indices] : use_list_) {
      if (infer_indices.empty())
        continue;

897
      // Union all infer_list_ indices that share the same Buffer object
898
899
900
901
902
      int first_idx = infer_indices[0];
      for (size_t i = 1; i < infer_indices.size(); i++) {
        uf.Union(first_idx, infer_indices[i]);
      }
    }
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
    // Additionally, union across buffers that share the same underlying
    // buffer->data (Var). This handles cases like reshape where multiple
    // Buffer objects alias the same storage.
    for (const auto &[var, buffers] : buffer_data_to_buffers_) {
      std::vector<int> merged;
      for (const auto &buf : buffers) {
        auto it = use_list_.find(buf);
        if (it != use_list_.end()) {
          const auto &vec = it->second;
          merged.insert(merged.end(), vec.begin(), vec.end());
        }
      }
      if (merged.size() > 1) {
        std::sort(merged.begin(), merged.end());
        merged.erase(std::unique(merged.begin(), merged.end()), merged.end());
        int first = merged[0];
        for (size_t i = 1; i < merged.size(); ++i) {
          uf.Union(first, merged[i]);
        }
      }
    }
924

925
926
927
928
929
    std::unordered_map<int, std::vector<int>> components;
    for (int i = 0; i < infer_list_.size(); i++) {
      int root = uf.Find(i);
      components[root].push_back(i);
    }
930
    // Create a map from root to buffers
931
932
933
934
935
    std::unordered_map<int, std::vector<Buffer>> components_buffers;
    for (const auto &[buffer, infer_indices] : use_list_) {
      int root = uf.Find(infer_indices[0]);
      components_buffers[root].push_back(buffer);
    }
936
937
    // Keep components_buffers for debug purpose
    (void)components_buffers;
938
939
940

    // For each component, try each op as root, and determine the least
    // replicated one
941
    std::deque<int> q;
942
    std::vector<bool> in_queue(infer_list_.size(), false);
943

944
    for (auto &&[root, members] : components) {
945
946
      DLOG(INFO) << "======================= processing component " << root
                 << '\n';
947
948
949
      decltype(infer_list_) best_infer_list;
      LayoutMap best_layout_map;
      int64_t min_reg_num = INT64_MAX;
950
      int min_reg_num_infer_root = -1;
951

952
      // Try each member as the root of inference for this component
953
      for (int attempt_infer_root : members) {
954
        DLOG(INFO) << "----------------------- try root " << attempt_infer_root
955
                   << " members " << members.size() << '\n';
956
        // Backup the current infer_list_ state
957
        auto back_infer_list = BackupInferList();
958
        // Copy the current layout_map for temporary use
959
960
961
        LayoutMap tmp_layout_map = layout_map;
        bool do_update = true;
        try {
962
          // Run inference starting from attempt_infer_root
963
964
965
966
          RunInferStep(attempt_infer_root, InferLevel::kFree, true,
                       tmp_layout_map, strict_layout_map, q, in_queue);
          FinishInferQueue(InferLevel::kFree, tmp_layout_map, strict_layout_map,
                           q, in_queue);
967
968
969

          // After the first search, run inference for all other members in
          // order
970
971
972
973
974
975
976
977
          for (int other_infer_root : members) {
            if (other_infer_root != attempt_infer_root) {
              RunInferStep(other_infer_root, InferLevel::kFree, true,
                           tmp_layout_map, strict_layout_map, q, in_queue);
              FinishInferQueue(InferLevel::kFree, tmp_layout_map,
                               strict_layout_map, q, in_queue);
            }
          }
978
        } catch (const LayoutConflictException &e) {
979
          do_update = false;
980
981
982
          DLOG(INFO) << "attempt failed due to LayoutConflictException "
                     << e.what() << '\n';
        } catch (const NormalizeIterException &e) {
983
          do_update = false;
984
985
          DLOG(INFO) << "attempt failed due to NormalizeIterException "
                     << e.what() << '\n';
986
987
988
989
        } catch (const LoopLayoutInjectiveException &e) {
          do_update = false;
          DLOG(INFO) << "attempt failed due to LoopLayoutInjectiveException "
                     << e.what() << '\n';
990
991
992
        }

        if (do_update) {
993
          // Compute the total register number for this layout
994
          int64_t reg_num = 0;
995
          for (const auto &[buffer, layout] : tmp_layout_map) {
996
997
998
999
            if (auto frag = layout.as<Fragment>()) {
              int64_t frag_reg_num = 1;
              for (auto i : frag.value()->OutputShape()) {
                auto pci = as_const_int(i);
1000
1001
1002
1003
1004
1005
1006
                ICHECK(pci != nullptr)
                    << "Can not use non-constant range to "
                       "iterate over a fragment/local "
                       "buffer. Non-constant shape expr is: "
                    << i
                    << ". This is possibly because you use symbolic shape when "
                       "accessing a fragment/local buffer.";
1007
1008
1009
1010
1011
                frag_reg_num *= *pci;
              }
              reg_num += frag_reg_num;
            }
          }
1012
          // Update the best plan if this one uses fewer registers
1013
1014
1015
          if (reg_num < min_reg_num ||
              (reg_num == min_reg_num &&
               attempt_infer_root < min_reg_num_infer_root)) {
1016
1017
            best_infer_list =
                BackupInferList(); // Use backup to avoid moving out infer_list_
1018
1019
            best_layout_map = tmp_layout_map;
            min_reg_num = reg_num;
1020
            min_reg_num_infer_root = attempt_infer_root;
1021
1022
          }
        }
1023
        // Restore infer_list_ state for the next attempt
1024
1025
        infer_list_ = std::move(back_infer_list);
      }
1026
1027
1028
1029
1030
1031
      ICHECK(min_reg_num < INT64_MAX) << "no available layout found" << '\n';
      // Apply the best plan for this component
      infer_list_ = std::move(best_infer_list);
      layout_map = best_layout_map;
      DLOG(INFO) << "[InferInFreeMode] Final selection is attempt_infer_root = "
                 << min_reg_num_infer_root << '\n';
1032
1033
    }
  }
1034
1035
1036
};

class LayoutInferencer : public IRMutatorWithAnalyzer {
1037
public:
1038
  static PrimFunc Substitute(PrimFunc f, bool skip_thread_partition = false) {
1039
    arith::Analyzer analyzer;
1040
    PrimFuncNode *fptr = f.CopyOnWrite();
1041
    fptr->body = ParallelLoopFuser::Fuse(f->body);
1042
    BufferUseDefCollector collector(skip_thread_partition);
1043
1044
    collector.Collect(f);
    auto result = collector.Run();
1045
    LayoutInferencer substituter(result, skip_thread_partition, &analyzer);
1046
1047
1048
1049
    fptr->body = substituter.VisitStmt(f->body);
    return f;
  }

1050
private:
1051
  LayoutInferencer(const LayoutInferenceResult &result,
1052
1053
                   bool skip_thread_partition, arith::Analyzer *analyzer)
      : arith::IRMutatorWithAnalyzer(analyzer), result_(result),
1054
        skip_thread_partition_(skip_thread_partition) {};
1055

1056
1057
  using arith::IRMutatorWithAnalyzer::IRMutatorWithAnalyzer;

1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
  /**
   * @brief Visit and mutate a Block node to attach inferred layout information.
   *
   * Converts the visited Block via the base visitor, asserts that every buffer
   * allocated with scope "local.framgent" has an inferred layout in
   * result_.layout_map, and attaches result_.layout_map to the Block's
   * annotations under attr::kLayoutMap.
   *
   * If any "local.framgent" buffer lacks an entry in result_.layout_map an
   * ICHECK will fail with the offending buffer printed.
   *
   * @return Stmt The (possibly modified) Block statement with the layout-map
   * annotation set.
   */
1072
  Stmt VisitStmt_(const BlockNode *op) final {
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
    Block block = Downcast<Block>(IRMutatorWithAnalyzer::VisitStmt_(op));

    for (auto buffer : block->alloc_buffers) {
      if (buffer.scope() == "local.framgent") {
        ICHECK(result_.layout_map.count(buffer))
            << "Cannot inference fragment layout for " << buffer;
      }
    }
    auto block_ptr = block.CopyOnWrite();
    block_ptr->annotations.Set(attr::kLayoutMap, result_.layout_map);
    return block;
  }

1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
  /**
   * @brief Visit and transform For nodes according to inferred layout
   * information.
   *
   * If the For node is present in result_.for_map, this method applies
   * loop-level layout-driven transformations: it optionally partitions the loop
   * across the thread index, vectorizes the loop body, and wraps the loop with
   * a predicate if one was inferred for the loop root.
   *
   * Detailed behavior:
   * - Reads reducer information from the For node's attr::kReducerInfo
   * annotation (if present) to detect reduction targets.
   * - Detects register-local buffer stores (buffers with scope "local") in the
   *   original loop body; if only register-local stores are present the loop is
   *   treated as a register-local scenario and is not partitioned across
   * threads.
   * - Obtains the loop layout from result_.for_map[root] and, unless the loop
   * is register-local or skip_thread_partition_ is set, partitions the loop via
   *   PartitionLoop using thread_var_ and analyzer_.
   * - Scans the transformed loop body to determine whether it accesses any
   *   non-local buffers (scopes other than "local" or "local.fragment").
   * - Scans the transformed loop body to detect reducers (based on
   * reducer_info). If a reducer is present the loop is NOT vectorized
   * (reduction axes are excluded from vectorization as a conservative
   * workaround).
   * - If the loop has non-local accesses and no reducer, the loop is vectorized
   *   via VectorizeLoop.
   * - If a predicate exists in result_.predicate_map for the loop root and the
   *   loop was partitioned, the method returns an IfThenElse surrounding the
   *   (possibly partitioned/vectorized) loop with that predicate; otherwise it
   *   returns the transformed For.
   *
   * @return The possibly transformed For statement (or an IfThenElse wrapping
   * it)
   */
1121
  Stmt VisitStmt_(const ForNode *op) final {
1122
1123
1124
1125
1126
    Map<Var, ReducerInfo> reducer_info;
    if (op->annotations.count(attr::kReducerInfo))
      reducer_info = op->annotations.Get(attr::kReducerInfo)
                         ->as<Map<Var, ReducerInfo>>()
                         .value();
1127
1128
1129
1130
1131
1132
    if (!result_.for_map.count(tvm::ffi::GetRef<For>(op))) {
      return IRMutatorWithAnalyzer::VisitStmt_(op);
    }
    // the analyzer will be modified in PartitionLoop and VectorizeLoop
    // we need to save its state to prevent conflicted bindings
    auto saved_analyzer = analyzer_->Clone();
1133
    For for_node = Downcast<For>(IRMutatorWithAnalyzer::VisitStmt_(op));
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
    auto root = tvm::ffi::GetRef<For>(op);
    // This check is a workaround to support T.Parallel for local buffers.
    // For example:
    //   for i in T.Parallel(1024):
    //     A_local[i] = A_global[i]
    // Here, A_local is a register-local buffer held independently by each
    // thread, so explicit thread binding is not required.
    bool store_into_local = false;
    PostOrderVisit(root, [&](const ObjectRef &obj) {
      if (const auto *store = obj.as<BufferStoreNode>()) {
        if (store->buffer.scope() == "local") {
          store_into_local = true;
1146
        }
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
        // if the case is like:
        // for i in T.Parallel(1024):
        //     A_local[i] = B_global[i]
        //     A_frag[i] = A_global[i]
        // exception will be raise in Parallel::LayoutInference
      }
    });
    // This check if for the loop that only manuplates "local" buffers,
    // for i in T.Parallel(1024):
    //     A_local[i] = B_local[i]
    // Though this might be illegal
    // We use PostOrderVisit to detect whether the loop only manuplates
    // "local" buffers, which indicates register usage and justifies skipping
    // thread binding.
    bool local_register_only = true;
    PostOrderVisit(root, [&](const ObjectRef &obj) {
      if (const auto *store = obj.as<BufferStoreNode>()) {
        if (store->buffer.scope() != "local") {
          local_register_only = false;
1166
        }
1167
1168
1169
1170
1171
1172
      } else if (const auto *load = obj.as<BufferLoadNode>()) {
        if (load->buffer.scope() != "local") {
          local_register_only = false;
        }
      }
    });
1173

1174
1175
1176
1177
1178
    auto loop_layout = result_.for_map[root];
    // FIXME: tell in-Parallel and out-of-Parallel `local`s apart
    // NOTE(lei): a bit ugly, we should rethink about this part in future.
    bool parallel_loop =
        !skip_thread_partition_ && !local_register_only && !store_into_local;
1179

1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
    if (parallel_loop) {
      for_node =
          PartitionLoop(for_node, thread_var_->var, analyzer_, loop_layout);
    }
    // If none thread bindings are provided, partition the loop
    bool has_non_local = false;
    PostOrderVisit(for_node->body, [&](const ObjectRef &obj) {
      if (const auto *load = obj.as<BufferLoadNode>()) {
        String scope = load->buffer.scope();
        if (scope != "local" && scope != "local.fragment") {
          has_non_local = true;
        }
      } else if (const auto *store = obj.as<BufferStoreNode>()) {
        String scope = store->buffer.scope();
        if (scope != "local" && scope != "local.fragment") {
          has_non_local = true;
        }
1197
      }
1198
1199
1200
1201
1202
1203
1204
1205
    });
    // Workaround: if reducer is presented, don't vectorize loop
    // Best solution should be isolate reduction axis out of vectorization
    bool has_reducer = false;
    PostOrderVisit(for_node->body, [&](const ObjectRef &obj) {
      if (!has_reducer)
        if (const auto *store = obj.as<BufferStoreNode>()) {
          has_reducer = reducer_info.count(store->buffer->data) != 0;
1206
        }
1207
1208
1209
1210
1211
1212
1213
1214
1215
    });

    // If a cast operation exists, vectorization may still be required
    bool has_cast_operations = false;
    PostOrderVisit(for_node->body, [&](const ObjectRef &obj) {
      if (const auto *cast = obj.as<CastNode>()) {
        // Check if this is a non-reducer store with Cast operation
        DataType src_type = cast->value.dtype();
        DataType dst_type = cast->dtype;
1216
1217
1218
1219
        bool src_ok =
            src_type.is_float() || src_type.is_bfloat() || src_type.is_float8();
        bool dst_ok =
            dst_type.is_float() || dst_type.is_bfloat() || dst_type.is_float8();
1220
1221
        if (src_ok && dst_ok && TargetIsCuda(Target::Current())) {
          has_cast_operations = true;
1222
        }
1223
      }
1224
    });
1225

1226
    if ((has_non_local || has_cast_operations) && !has_reducer) {
1227
      DLOG(INFO) << "Try to vectorize loop";
1228
1229
1230
1231
1232
1233
1234
      for_node = VectorizeLoop(for_node, saved_analyzer.get());
    }

    if (result_.predicate_map.count(root) && parallel_loop) {
      return IfThenElse(result_.predicate_map[root], for_node);
    } else {
      return for_node;
1235
1236
1237
    }
  }

1238
  Stmt VisitStmt_(const AttrStmtNode *op) final {
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
    if (op->attr_key == tir::attr::thread_extent) {
      IterVar iv = Downcast<IterVar>(op->node);
      ICHECK_NE(iv->thread_tag.length(), 0U);
      if (iv->thread_tag == "threadIdx.x") {
        thread_var_ = iv;
      }
    }
    return IRMutatorWithAnalyzer::VisitStmt_(op);
  }

1249
private:
1250
  const LayoutInferenceResult result_;
1251
1252
  IterVar thread_var_ = IterVar(Range::FromMinExtent(0, 1), Var("v_thread"),
                                IterVarType::kDataPar);
1253
  bool skip_thread_partition_{false};
1254
1255
1256
1257
};

tvm::transform::Pass LayoutInference() {
  using namespace tir::transform;
1258
  auto pass_func = [=](PrimFunc f, const IRModule &m, const PassContext &ctx) {
1259
    f.CopyOnWrite()->body = ParallelLoopTransformer::Substitute(f->body);
1260
    ThreadBindingCollector collector;
1261
    collector(f->body);
1262
    bool has_thread_binding = !collector.thread_binding_.empty();
1263
    bool skip_thread_partition = !has_thread_binding;
1264
    return LayoutInferencer::Substitute(std::move(f), skip_thread_partition);
1265
1266
1267
1268
  };
  return CreatePrimFuncPass(pass_func, 0, "tl.LayoutInference", {});
}

1269
TVM_FFI_STATIC_INIT_BLOCK() {
1270
1271
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def("tl.transform.LayoutInference", LayoutInference);
1272
}
1273

1274
1275
} // namespace tl
} // namespace tvm