layout_inference.cc 41.8 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 <memory>
16
17
#include <queue>

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

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

namespace tvm {
namespace tl {

37
38
39
using namespace tir;

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

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

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

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

66
class BufferUseDefCollector : public IRVisitorWithAnalyzer {
67
public:
68
69
  BufferUseDefCollector(bool skip_thread_partition)
      : skip_thread_partition_(skip_thread_partition) {}
70

71
72
  using arith::IRVisitorWithAnalyzer::IRVisitorWithAnalyzer;

73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
  void RunInferStep(int cur_infer_id, InferLevel level, bool update_queue,
                    LayoutMap &layout_map, const LayoutMap &strict_layout_map,
                    std::queue<int> &q, std::vector<bool> &in_queue) {
    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];
89
    arith::Analyzer *cur_analyzer = analyzer_vec_[cur_infer_id].get();
90
    auto buffer_oob = buffer_oob_vec_[cur_infer_id];
91
    // Double-check that 'next' is valid
92
93
    ICHECK(next.defined()) << "infer_list_[" << cur_infer_id
                           << "] is null inside run_infer_step.";
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110

    // 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
111
112
    auto updates =
        next->InferLayout(LayoutInferArgs{target_, thread_bounds, layout_map,
113
                                          cur_analyzer, buffer_oob},
114
                          level);
115
116
    // Process the returned updates
    for (const auto &[buffer, layout] : updates) {
117

118
119
120
121
      // Basic validity checks
      ICHECK(buffer.defined()) << "InferLayout returned an undefined buffer.";
      ICHECK(layout.defined()) << "InferLayout returned an undefined layout.";

122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
      // 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 =
              shapes_equal ? src_layout
                           : src_layout->Reshape(sib->shape, &analyzer_);
          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]) {
                if (!in_queue[idx] && idx != cur_infer_id) {
                  in_queue[idx] = true;
                  q.push(idx);
                }
              }
            }
          }
        }
      };

165
166
167
168
169
170
      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
171
          // `TileOperator`s.
172
173
174
175
176

          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();
177
          const auto &dst_layout = dst_layout_opt.value();
178
179
180
181
182
          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();
183
          const auto &src_layout = src_layout_opt.value();
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
          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);
199
200
            // Propagate to alias buffers as well
            propagate_alias(buffer, layout);
201
202
203
204
            continue;
          }
        }
        // If already in map, ensure they are structurally equal
205
        ICHECK(layout->IsEqual(layout_map[buffer].get()))
206
207
208
            << "Get different layout for " << buffer
            << "\n current layout: " << layout->DebugOutput()
            << "\n previous layout: " << layout_map[buffer]->DebugOutput();
209
210
        // Ensure aliases are consistent too
        propagate_alias(buffer, layout);
211
212
213
      } else {
        // Otherwise, update map
        layout_map.Set(buffer, layout);
214
215
        // Propagate to alias buffers (may enqueue their users)
        propagate_alias(buffer, layout);
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
        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 << ".";

          if (!in_queue[idx] && idx != cur_infer_id) {
            in_queue[idx] = true;
            q.push(idx);
          }
        }
      }
    }
  };

  void FinishInferQueue(InferLevel level, LayoutMap &layout_map,
                        const LayoutMap &strict_layout_map, std::queue<int> &q,
                        std::vector<bool> &in_queue) {
    auto num_infer = infer_list_.size();
    while (!q.empty()) {
      int cur_infer_id = q.front();
      q.pop();
      // 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);
    }
  };

262
  LayoutInferenceResult Run() {
263
264
265
266
267
    // 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.";
268
269
270
    ICHECK_EQ(thread_bounds_vec_.size(), infer_list_.size())
        << "Size mismatch: thread_bounds_vec_ and infer_list_ must match in "
           "length.";
271
272
273
    ICHECK_EQ(analyzer_vec_.size(), infer_list_.size())
        << "Size mismatch: analyzer_vec_ and infer_list_ must match in "
           "length.";
274
275
276
    ICHECK_EQ(buffer_oob_vec_.size(), infer_list_.size())
        << "Size mismatch: buffer_oob_vec_ and infer_list_ must match in "
           "length.";
277

278
279
280
281
282
    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';
    }

283
284
285
286
    // 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
287
    Map<Buffer, Layout> layout_map = annotated_layout_map_;
288
    Map<Buffer, Layout> strict_layout_map;
289
290
    int num_infer = infer_list_.size();

291
    // Prepare BFS queue for iterative inference
292
293
    std::queue<int> q;
    std::vector<bool> in_queue(num_infer, true);
294
295
    for (int i = 0; i < num_infer; i++) {
      // Check that each infer_list_ entry is valid
296
      ICHECK(infer_list_[i].defined())
297
298
299
300
301
302
303
          << "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_;
      }
304
      q.push(i);
305
    }
306

307
    // step 1: infer strict layout
308
    for (int i = 0; i < num_infer; i++) {
309
310
      RunInferStep(i, InferLevel::kStrict, false, layout_map, strict_layout_map,
                   q, in_queue);
311
312
    }

313
314
315
316
    for (const auto &[buffer, layout] : layout_map) {
      strict_layout_map.Set(buffer, layout);
    }

317
    // step 2: infer common layout with BFS
318
319
    FinishInferQueue(InferLevel::kCommon, layout_map, strict_layout_map, q,
                     in_queue);
320

321
    // step 3: relax constraints to free and re-run
322
323
    InferInFreeMode(layout_map, strict_layout_map);

324
325
326
327
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
359
360
361
362
363
    // 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;
              }
            }
          }

          Layout reshaped =
              shapes_equal
                  ? rep_layout.value()
                  : rep_layout.value()->Reshape(buf->shape, &analyzer_);
          layout_map.Set(buf, reshaped);
        }
      }
    }

364
    // Check that all local.fragment buffers have inferred layouts
365
    for (const auto &[buffer, _] : use_list_) {
366
367
368
369
370
      if (buffer.scope() == "local.fragment") {
        ICHECK_NE(layout_map.count(buffer), 0)
            << "The layout for fragment " << buffer
            << " can not be inferred correctly.";
      }
371
372
    }

373
    // Collect layout info for For nodes
374
375
    Map<For, Fragment> for_map;
    Map<For, PrimExpr> predicate_map;
376
377
378
    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++) {
379
      TileOperator base_infer = std::move(infer_list_[i]);
380
381
      auto thread_var = thread_var_vec_[i];

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

395
        if (auto predicate = for_infer->GetPredicate(thread_var->var)) {
396
          predicate_map.Set(for_infer->GetRoot(), predicate.value());
397
        }
398
399
400
401
402
403
      }
    }

    return {layout_map, for_map, predicate_map};
  }

404
405
  void Collect(const PrimFunc &f) {
    for (const auto &[_, buffer] : f->buffer_map) {
406
407
408
409
410
411
412
      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});
      }
413
414
    }
    auto target = f->GetAttr<Target>(tvm::attr::kTarget);
415
416
    ICHECK(target.defined())
        << "Layout_Inference: Require the target attribute";
417
418
419
420
    target_ = target.value();
    this->operator()(f->body);
  }

421
private:
422
423
424
425
426
427
428
429
430
431
432
433
  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;
  }

434
  void VisitExpr_(const CallNode *op) final {
435
    IRVisitorWithAnalyzer::VisitExpr_(op);
436
    // Do not analysis the call node to the global function.
437
438
    if (op->op.as<GlobalVarNode>())
      return;
439

440
    auto p = ParseOperator(tvm::ffi::GetRef<Call>(op));
441
    if (p.defined()) {
442
      for (const auto &arg : op->args) {
443
444
        if (auto buffer = getBufferFromAccessPtr(arg)) {
          addToUseList(buffer.value());
445
446
        } else if (auto buffer = getBufferFromRegion(arg)) {
          addToUseList(buffer.value());
447
448
        }
      }
449
      // Compute thread_var_ and thread_bounds_
450
      thread_var_vec_.push_back(thread_var_);
451
452
453
454
      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;
455
        auto extent = max_value - min_value + 1;
456
457
        auto dtype = thread_var_->var.dtype();
        thread_bounds_vec_.push_back(Range::FromMinExtent(
458
            IntImm(dtype, min_value), IntImm(dtype, extent)));
459
460
461
      } else {
        thread_bounds_vec_.push_back(Range::FromMinExtent(0, 1));
      }
462
      analyzer_vec_.push_back(analyzer_.Clone());
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

      // 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_
494
      infer_list_stmt_.push_back(tvm::ffi::GetRef<ObjectRef>(op));
495
      infer_list_.push_back(std::move(p));
496
497
498
    }
  }

499
  Optional<Buffer> getBufferFromAccessPtr(const PrimExpr &expr) {
500
501
502
    if (auto bl = expr.as<BufferLoadNode>()) {
      return bl->buffer;
    }
503
    auto call = expr.as<CallNode>();
504
505
506
507
    if (!call) {
      return std::nullopt;
    }
    if (call->op.same_as(builtin::tvm_access_ptr())) {
508
509
      auto var_opt = call->args[1].as<Var>();
      if (!var_opt.has_value()) {
510
511
        LOG(WARNING) << "[getBufferFromAccessPtr] args[1] is not a Var, type: "
                     << call->args[1]->GetTypeKey();
512
513
        return std::nullopt;
      }
514
      const auto &var = var_opt.value();
515
516
517
518
519
520
521
      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;
522
523
524
525
526
527
528
529
530
531
532
533
    }
    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;
      }
534
    }
535
    return std::nullopt;
536
537
  }

538
  void addToUseList(const Buffer &buffer) {
539
540
541
542
543
544
545
    int infer_idx = infer_list_.size();
    if (use_list_.find(buffer) == use_list_.end()) {
      use_list_[buffer] = {};
    }
    use_list_[buffer].push_back(infer_idx);
  }

546
  void VisitStmt_(const ForNode *op) final {
547
    if (op->kind == ForKind::kParallel) {
548
      auto infer = ParallelOp(tvm::ffi::GetRef<For>(op));
549
      for (const auto &[buffer, _] : infer->GetIndiceMap()) {
550
551
        addToUseList(buffer);
      }
552
      infer_list_stmt_.push_back(tvm::ffi::GetRef<ObjectRef>(op));
553
554
      infer_list_.push_back(std::move(infer));
      thread_var_vec_.push_back(thread_var_);
555
556
557
558
      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();
559
560
        auto extent =
            const_int_bound->max_value - const_int_bound->min_value + 1;
561
        thread_bounds_vec_.push_back(Range::FromMinExtent(
562
            IntImm(dtype, const_int_bound->min_value), IntImm(dtype, extent)));
563
564
565
      } else {
        thread_bounds_vec_.push_back(Range::FromMinExtent(0, 1));
      }
566
      analyzer_vec_.push_back(analyzer_.Clone());
567
      buffer_oob_vec_.push_back(false);
568
    } else {
569
      IRVisitorWithAnalyzer::VisitStmt(op->body);
570
571
572
    }
  }

573
  void VisitStmt_(const BlockNode *op) final {
574
    for (auto buffer : op->alloc_buffers) {
575
576
577
578
579
580
581
      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});
      }
582
    }
583
584
585
586
587
588

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

    // After visiting, apply layouts to all collected buffers
589
    if (op->annotations.count(attr::kLayoutMap)) {
590
      // Check if the layout map is Map<Var, Layout>
591
592
593
      auto map =
          op->annotations.Get(attr::kLayoutMap)->as<Map<Var, Layout>>().value();
      for (const auto &[var, layout] : map) {
594
        ICHECK(buffer_data_to_buffers_.count(var))
595
            << "buffer " << var << " is not found in the block";
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
        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 {
            auto reshaped_layout = layout->Reshape(buffer->shape, &analyzer_);
            annotated_layout_map_.Set(buffer, reshaped_layout);
          }
        }
622
623
624
625
      }
    }
  }

626
  void VisitStmt_(const AttrStmtNode *op) final {
627
628
629
630
631
632
633
    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;
      }
    }
634
    IRVisitorWithAnalyzer::VisitStmt_(op);
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
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
  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_;
698
  std::vector<ObjectRef> infer_list_stmt_;
699
  std::vector<TileOperator> infer_list_;
700
701
  std::unordered_map<Buffer, std::vector<int>, ObjectPtrHash, ObjectPtrEqual>
      use_list_;
702
703
704
705
  // 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);
706
  std::vector<IterVar> thread_var_vec_;
707
  std::vector<Range> thread_bounds_vec_;
708
  std::vector<std::unique_ptr<arith::Analyzer>> analyzer_vec_;
709
  std::vector<bool> buffer_oob_vec_;
710
711
  Target target_;
  LayoutMap annotated_layout_map_;
712
  bool skip_thread_partition_{false};
713

714
715
  std::vector<TileOperator> BackupInferList() {
    std::vector<TileOperator> back_infer_list;
716
717
718
719
720
721
722
723
724
    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) {
725
726
727
728
729
730
731

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

732
733
734
735
736
737
738
739
740
    // 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;

741
      // Union all infer_list_ indices that share the same Buffer object
742
743
744
745
746
      int first_idx = infer_indices[0];
      for (size_t i = 1; i < infer_indices.size(); i++) {
        uf.Union(first_idx, infer_indices[i]);
      }
    }
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
    // 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]);
        }
      }
    }
768
769
770
771
772
    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);
    }
773
    // Create a map from root to buffers
774
775
776
777
778
    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);
    }
779
780
    // Keep components_buffers for debug purpose
    (void)components_buffers;
781
782
783
784
785

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

787
    for (auto &&[root, members] : components) {
788
789
      DLOG(INFO) << "======================= processing component " << root
                 << '\n';
790
791
792
      decltype(infer_list_) best_infer_list;
      LayoutMap best_layout_map;
      int64_t min_reg_num = INT64_MAX;
793
      int min_reg_num_infer_root = -1;
794

795
      // Try each member as the root of inference for this component
796
      for (int attempt_infer_root : members) {
797
798
799
        DLOG(INFO) << "----------------------- try root " << attempt_infer_root
                   << '\n';
        // Backup the current infer_list_ state
800
        auto back_infer_list = BackupInferList();
801
        // Copy the current layout_map for temporary use
802
803
804
        LayoutMap tmp_layout_map = layout_map;
        bool do_update = true;
        try {
805
          // Run inference starting from attempt_infer_root
806
807
808
809
          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);
810
811
812

          // After the first search, run inference for all other members in
          // order
813
814
815
816
817
818
819
820
          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);
            }
          }
821
        } catch (const LayoutConflictException &e) {
822
          do_update = false;
823
824
825
          DLOG(INFO) << "attempt failed due to LayoutConflictException "
                     << e.what() << '\n';
        } catch (const NormalizeIterException &e) {
826
          do_update = false;
827
828
          DLOG(INFO) << "attempt failed due to NormalizeIterException "
                     << e.what() << '\n';
829
830
831
        }

        if (do_update) {
832
          // Compute the total register number for this layout
833
          int64_t reg_num = 0;
834
          for (const auto &[buffer, layout] : tmp_layout_map) {
835
836
837
838
            if (auto frag = layout.as<Fragment>()) {
              int64_t frag_reg_num = 1;
              for (auto i : frag.value()->OutputShape()) {
                auto pci = as_const_int(i);
839
840
841
842
843
844
845
                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.";
846
847
848
849
850
                frag_reg_num *= *pci;
              }
              reg_num += frag_reg_num;
            }
          }
851
          // Update the best plan if this one uses fewer registers
852
853
854
          if (reg_num < min_reg_num ||
              (reg_num == min_reg_num &&
               attempt_infer_root < min_reg_num_infer_root)) {
855
856
            best_infer_list =
                BackupInferList(); // Use backup to avoid moving out infer_list_
857
858
            best_layout_map = tmp_layout_map;
            min_reg_num = reg_num;
859
            min_reg_num_infer_root = attempt_infer_root;
860
861
          }
        }
862
        // Restore infer_list_ state for the next attempt
863
864
        infer_list_ = std::move(back_infer_list);
      }
865
866
867
868
869
870
      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';
871
872
    }
  }
873
874
875
};

class LayoutInferencer : public IRMutatorWithAnalyzer {
876
public:
877
  static PrimFunc Substitute(PrimFunc f, bool skip_thread_partition = false) {
878
    arith::Analyzer analyzer;
879
    PrimFuncNode *fptr = f.CopyOnWrite();
880
    fptr->body = ParallelLoopFuser::Fuse(f->body);
881
    BufferUseDefCollector collector(skip_thread_partition);
882
883
    collector.Collect(f);
    auto result = collector.Run();
884
    LayoutInferencer substituter(result, skip_thread_partition, &analyzer);
885
886
887
888
    fptr->body = substituter.VisitStmt(f->body);
    return f;
  }

889
private:
890
  LayoutInferencer(const LayoutInferenceResult &result,
891
892
                   bool skip_thread_partition, arith::Analyzer *analyzer)
      : arith::IRMutatorWithAnalyzer(analyzer), result_(result),
893
        skip_thread_partition_(skip_thread_partition) {};
894

895
896
  using arith::IRMutatorWithAnalyzer::IRMutatorWithAnalyzer;

897
898
899
900
901
902
903
904
905
906
907
908
909
910
  /**
   * @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.
   */
911
  Stmt VisitStmt_(const BlockNode *op) final {
912
913
914
915
916
917
918
919
920
921
922
923
924
    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;
  }

925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
  /**
   * @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)
   */
960
  Stmt VisitStmt_(const ForNode *op) final {
961
962
963
964
965
966
    Map<Var, ReducerInfo> reducer_info;
    if (op->annotations.count(attr::kReducerInfo))
      reducer_info = op->annotations.Get(attr::kReducerInfo)
                         ->as<Map<Var, ReducerInfo>>()
                         .value();

967
    For for_node = Downcast<For>(IRMutatorWithAnalyzer::VisitStmt_(op));
968
969
    if (result_.for_map.count(tvm::ffi::GetRef<For>(op))) {
      auto root = tvm::ffi::GetRef<For>(op);
970
971
972
973
974
975
      // 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.
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
      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;
          }
          // 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
993
994
      // We use PostOrderVisit to detect whether the loop only manuplates
      // "local" buffers, which indicates register usage and justifies skipping
995
      // thread binding.
996
      bool local_register_only = true;
997
998
      PostOrderVisit(root, [&](const ObjectRef &obj) {
        if (const auto *store = obj.as<BufferStoreNode>()) {
999
1000
1001
1002
1003
1004
          if (store->buffer.scope() != "local") {
            local_register_only = false;
          }
        } else if (const auto *load = obj.as<BufferLoadNode>()) {
          if (load->buffer.scope() != "local") {
            local_register_only = false;
1005
1006
1007
1008
          }
        }
      });

1009
      auto loop_layout = result_.for_map[root];
1010
      // FIXME: tell in-Parallel and out-of-Parallel `local`s apart
1011
1012
1013
      // 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;
1014

1015
      if (parallel_loop) {
1016
1017
1018
        for_node =
            PartitionLoop(for_node, thread_var_->var, analyzer_, loop_layout);
      }
1019
      // If none thread bindings are provided, partition the loop
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
      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;
          }
        }
      });
1034
1035
1036
1037
1038
1039
1040
1041
1042
      // 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;
          }
      });
1043

1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
      // 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 *store = obj.as<BufferStoreNode>()) {
          // Check if this is a non-reducer store with Cast operation
          if (store->value.as<CastNode>()) {
            has_cast_operations = true;
          }
        }
      });

      if ((has_non_local || has_cast_operations) && !has_reducer) {
1056
        for_node = VectorizeLoop(for_node, analyzer_);
1057
      }
1058

1059
1060
      if (result_.predicate_map.count(root) && parallel_loop) {
        return IfThenElse(result_.predicate_map[root], for_node);
1061
1062
1063
1064
1065
1066
1067
      } else {
        return for_node;
      }
    }
    return for_node;
  }

1068
  Stmt VisitStmt_(const AttrStmtNode *op) final {
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
    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);
  }

1079
private:
1080
  const LayoutInferenceResult result_;
1081
1082
  IterVar thread_var_ = IterVar(Range::FromMinExtent(0, 1), Var("v_thread"),
                                IterVarType::kDataPar);
1083
  bool skip_thread_partition_{false};
1084
1085
1086
1087
};

tvm::transform::Pass LayoutInference() {
  using namespace tir::transform;
1088
  auto pass_func = [=](PrimFunc f, const IRModule &m, const PassContext &ctx) {
1089
    f.CopyOnWrite()->body = ParallelLoopTransformer::Substitute(f->body);
1090
    ThreadBindingCollector collector;
1091
    collector(f->body);
1092
    bool has_thread_binding = !collector.thread_binding_.empty();
1093
    bool skip_thread_partition = !has_thread_binding;
1094
    return LayoutInferencer::Substitute(std::move(f), skip_thread_partition);
1095
1096
1097
1098
  };
  return CreatePrimFuncPass(pass_func, 0, "tl.LayoutInference", {});
}

1099
TVM_FFI_STATIC_INIT_BLOCK() {
1100
1101
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def("tl.transform.LayoutInference", LayoutInference);
1102
}
1103

1104
1105
} // namespace tl
} // namespace tvm