"src/target/codegen_hip.cc" did not exist on "64f17c2f369e612cc297d358f607307a615bbb59"
layout_inference.cc 40.7 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
16
#include <queue>

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

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

namespace tvm {
namespace tl {

36
37
38
using namespace tir;

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

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

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

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

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

70
71
  using arith::IRVisitorWithAnalyzer::IRVisitorWithAnalyzer;

72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
  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];
88
    auto buffer_oob = buffer_oob_vec_[cur_infer_id];
89
    // Double-check that 'next' is valid
90
91
    ICHECK(next.defined()) << "infer_list_[" << cur_infer_id
                           << "] is null inside run_infer_step.";
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108

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

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

120
121
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
      // 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);
                }
              }
            }
          }
        }
      };

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

          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();
175
          const auto &dst_layout = dst_layout_opt.value();
176
177
178
179
180
          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();
181
          const auto &src_layout = src_layout_opt.value();
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
          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);
197
198
            // Propagate to alias buffers as well
            propagate_alias(buffer, layout);
199
200
201
202
            continue;
          }
        }
        // If already in map, ensure they are structurally equal
203
        ICHECK(layout->IsEqual(layout_map[buffer].get()))
204
205
206
            << "Get different layout for " << buffer
            << "\n current layout: " << layout->DebugOutput()
            << "\n previous layout: " << layout_map[buffer]->DebugOutput();
207
208
        // Ensure aliases are consistent too
        propagate_alias(buffer, layout);
209
210
211
      } else {
        // Otherwise, update map
        layout_map.Set(buffer, layout);
212
213
        // Propagate to alias buffers (may enqueue their users)
        propagate_alias(buffer, layout);
214
215
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
        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);
    }
  };

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

273
274
275
276
277
    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';
    }

278
279
280
281
    // 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
282
    Map<Buffer, Layout> layout_map = annotated_layout_map_;
283
    Map<Buffer, Layout> strict_layout_map;
284
285
    int num_infer = infer_list_.size();

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

302
    // step 1: infer strict layout
303
    for (int i = 0; i < num_infer; i++) {
304
305
      RunInferStep(i, InferLevel::kStrict, false, layout_map, strict_layout_map,
                   q, in_queue);
306
307
    }

308
309
310
311
    for (const auto &[buffer, layout] : layout_map) {
      strict_layout_map.Set(buffer, layout);
    }

312
    // step 2: infer common layout with BFS
313
314
    FinishInferQueue(InferLevel::kCommon, layout_map, strict_layout_map, q,
                     in_queue);
315

316
    // step 3: relax constraints to free and re-run
317
318
    InferInFreeMode(layout_map, strict_layout_map);

319
320
321
322
323
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
    // 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);
        }
      }
    }

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

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

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

390
        if (auto predicate = for_infer->GetPredicate(thread_var->var)) {
391
          predicate_map.Set(for_infer->GetRoot(), predicate.value());
392
        }
393
394
395
396
397
398
      }
    }

    return {layout_map, for_map, predicate_map};
  }

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

416
private:
417
418
419
420
421
422
423
424
425
426
427
428
  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;
  }

429
  void VisitExpr_(const CallNode *op) final {
430
    IRVisitorWithAnalyzer::VisitExpr_(op);
431
    // Do not analysis the call node to the global function.
432
433
    if (op->op.as<GlobalVarNode>())
      return;
434

435
    auto p = ParseOperator(tvm::ffi::GetRef<Call>(op), GetBufferMap());
436
    if (p.defined()) {
437
      for (const auto &arg : op->args) {
438
439
440
441
        if (auto buffer = getBufferFromAccessPtr(arg)) {
          addToUseList(buffer.value());
        }
      }
442
      // Compute thread_var_ and thread_bounds_
443
      thread_var_vec_.push_back(thread_var_);
444
445
446
447
      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;
448
        auto extent = max_value - min_value + 1;
449
450
        auto dtype = thread_var_->var.dtype();
        thread_bounds_vec_.push_back(Range::FromMinExtent(
451
            IntImm(dtype, min_value), IntImm(dtype, extent)));
452
453
454
      } else {
        thread_bounds_vec_.push_back(Range::FromMinExtent(0, 1));
      }
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485

      // 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_
486
      infer_list_stmt_.push_back(tvm::ffi::GetRef<ObjectRef>(op));
487
      infer_list_.push_back(std::move(p));
488
489
490
    }
  }

491
  Optional<Buffer> getBufferFromAccessPtr(const PrimExpr &expr) {
492
    auto call = expr.as<CallNode>();
493
494
495
496
    if (!call) {
      return std::nullopt;
    }
    if (call->op.same_as(builtin::tvm_access_ptr())) {
497
498
      auto var_opt = call->args[1].as<Var>();
      if (!var_opt.has_value()) {
499
500
        LOG(WARNING) << "[getBufferFromAccessPtr] args[1] is not a Var, type: "
                     << call->args[1]->GetTypeKey();
501
502
        return std::nullopt;
      }
503
      const auto &var = var_opt.value();
504
505
506
507
508
509
510
      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;
511
512
    } else if (call->op.same_as(RegionOp::Get())) {
      return call->args[0].as<BufferLoadNode>()->buffer;
513
    }
514
    return std::nullopt;
515
516
  }

517
  void addToUseList(const Buffer &buffer) {
518
519
520
521
522
523
524
    int infer_idx = infer_list_.size();
    if (use_list_.find(buffer) == use_list_.end()) {
      use_list_[buffer] = {};
    }
    use_list_[buffer].push_back(infer_idx);
  }

525
  void VisitStmt_(const ForNode *op) final {
526
    if (op->kind == ForKind::kParallel) {
527
      auto infer = ParallelOp(tvm::ffi::GetRef<For>(op));
528
      for (const auto &[buffer, _] : infer->GetIndiceMap()) {
529
530
        addToUseList(buffer);
      }
531
      infer_list_stmt_.push_back(tvm::ffi::GetRef<ObjectRef>(op));
532
533
      infer_list_.push_back(std::move(infer));
      thread_var_vec_.push_back(thread_var_);
534
535
536
537
      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();
538
539
        auto extent =
            const_int_bound->max_value - const_int_bound->min_value + 1;
540
        thread_bounds_vec_.push_back(Range::FromMinExtent(
541
            IntImm(dtype, const_int_bound->min_value), IntImm(dtype, extent)));
542
543
544
      } else {
        thread_bounds_vec_.push_back(Range::FromMinExtent(0, 1));
      }
545
      buffer_oob_vec_.push_back(false);
546
    } else {
547
      IRVisitorWithAnalyzer::VisitStmt(op->body);
548
549
550
    }
  }

551
  void VisitStmt_(const BlockNode *op) final {
552
    for (auto buffer : op->alloc_buffers) {
553
554
555
556
557
558
559
      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});
      }
560
    }
561
562
563
564
565
566

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

    // After visiting, apply layouts to all collected buffers
567
    if (op->annotations.count(attr::kLayoutMap)) {
568
      // Check if the layout map is Map<Var, Layout>
569
570
571
      auto map =
          op->annotations.Get(attr::kLayoutMap)->as<Map<Var, Layout>>().value();
      for (const auto &[var, layout] : map) {
572
        ICHECK(buffer_data_to_buffers_.count(var))
573
            << "buffer " << var << " is not found in the block";
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
        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);
          }
        }
600
601
602
603
      }
    }
  }

604
  void VisitStmt_(const AttrStmtNode *op) final {
605
606
607
608
609
610
611
    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;
      }
    }
612
    IRVisitorWithAnalyzer::VisitStmt_(op);
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
  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_;
676
  std::vector<ObjectRef> infer_list_stmt_;
677
  std::vector<TileOperator> infer_list_;
678
679
  std::unordered_map<Buffer, std::vector<int>, ObjectPtrHash, ObjectPtrEqual>
      use_list_;
680
681
682
683
  // 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);
684
  std::vector<IterVar> thread_var_vec_;
685
  std::vector<Range> thread_bounds_vec_;
686
  std::vector<bool> buffer_oob_vec_;
687
688
  Target target_;
  LayoutMap annotated_layout_map_;
689
  bool skip_thread_partition_{false};
690

691
692
  std::vector<TileOperator> BackupInferList() {
    std::vector<TileOperator> back_infer_list;
693
694
695
696
697
698
699
700
701
    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) {
702
703
704
705
706
707
708

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

709
710
711
712
713
714
715
716
717
    // 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;

718
      // Union all infer_list_ indices that share the same Buffer object
719
720
721
722
723
      int first_idx = infer_indices[0];
      for (size_t i = 1; i < infer_indices.size(); i++) {
        uf.Union(first_idx, infer_indices[i]);
      }
    }
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
    // 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]);
        }
      }
    }
745
746
747
748
749
    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);
    }
750
    // Create a map from root to buffers
751
752
753
754
755
    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);
    }
756
757
    // Keep components_buffers for debug purpose
    (void)components_buffers;
758
759
760
761
762

    // 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);
763

764
    for (auto &&[root, members] : components) {
765
766
      DLOG(INFO) << "======================= processing component " << root
                 << '\n';
767
768
769
      decltype(infer_list_) best_infer_list;
      LayoutMap best_layout_map;
      int64_t min_reg_num = INT64_MAX;
770
      int min_reg_num_infer_root = -1;
771

772
      // Try each member as the root of inference for this component
773
      for (int attempt_infer_root : members) {
774
775
776
        DLOG(INFO) << "----------------------- try root " << attempt_infer_root
                   << '\n';
        // Backup the current infer_list_ state
777
        auto back_infer_list = BackupInferList();
778
        // Copy the current layout_map for temporary use
779
780
781
        LayoutMap tmp_layout_map = layout_map;
        bool do_update = true;
        try {
782
          // Run inference starting from attempt_infer_root
783
784
785
786
          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);
787
788
789

          // After the first search, run inference for all other members in
          // order
790
791
792
793
794
795
796
797
          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);
            }
          }
798
        } catch (const LayoutConflictException &e) {
799
          do_update = false;
800
801
802
          DLOG(INFO) << "attempt failed due to LayoutConflictException "
                     << e.what() << '\n';
        } catch (const NormalizeIterException &e) {
803
          do_update = false;
804
805
          DLOG(INFO) << "attempt failed due to NormalizeIterException "
                     << e.what() << '\n';
806
807
808
        }

        if (do_update) {
809
          // Compute the total register number for this layout
810
          int64_t reg_num = 0;
811
          for (const auto &[buffer, layout] : tmp_layout_map) {
812
813
814
815
816
817
818
819
820
821
            if (auto frag = layout.as<Fragment>()) {
              int64_t frag_reg_num = 1;
              for (auto i : frag.value()->OutputShape()) {
                auto pci = as_const_int(i);
                ICHECK(pci != nullptr);
                frag_reg_num *= *pci;
              }
              reg_num += frag_reg_num;
            }
          }
822
          // Update the best plan if this one uses fewer registers
823
824
825
          if (reg_num < min_reg_num ||
              (reg_num == min_reg_num &&
               attempt_infer_root < min_reg_num_infer_root)) {
826
827
            best_infer_list =
                BackupInferList(); // Use backup to avoid moving out infer_list_
828
829
            best_layout_map = tmp_layout_map;
            min_reg_num = reg_num;
830
            min_reg_num_infer_root = attempt_infer_root;
831
832
          }
        }
833
        // Restore infer_list_ state for the next attempt
834
835
        infer_list_ = std::move(back_infer_list);
      }
836
837
838
839
840
841
      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';
842
843
    }
  }
844
845
846
};

class LayoutInferencer : public IRMutatorWithAnalyzer {
847
public:
848
  static PrimFunc Substitute(PrimFunc f, bool skip_thread_partition = false) {
849
    arith::Analyzer analyzer;
850
    PrimFuncNode *fptr = f.CopyOnWrite();
851
    fptr->body = ParallelLoopFuser::Fuse(f->body);
852
    BufferUseDefCollector collector(skip_thread_partition);
853
854
    collector.Collect(f);
    auto result = collector.Run();
855
    LayoutInferencer substituter(result, skip_thread_partition, &analyzer);
856
857
858
859
    fptr->body = substituter.VisitStmt(f->body);
    return f;
  }

860
private:
861
  LayoutInferencer(const LayoutInferenceResult &result,
862
863
                   bool skip_thread_partition, arith::Analyzer *analyzer)
      : arith::IRMutatorWithAnalyzer(analyzer), result_(result),
864
        skip_thread_partition_(skip_thread_partition) {};
865

866
867
  using arith::IRMutatorWithAnalyzer::IRMutatorWithAnalyzer;

868
869
870
871
872
873
874
875
876
877
878
879
880
881
  /**
   * @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.
   */
882
  Stmt VisitStmt_(const BlockNode *op) final {
883
884
885
886
887
888
889
890
891
892
893
894
895
    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;
  }

896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
  /**
   * @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)
   */
931
  Stmt VisitStmt_(const ForNode *op) final {
932
933
934
935
936
937
    Map<Var, ReducerInfo> reducer_info;
    if (op->annotations.count(attr::kReducerInfo))
      reducer_info = op->annotations.Get(attr::kReducerInfo)
                         ->as<Map<Var, ReducerInfo>>()
                         .value();

938
    For for_node = Downcast<For>(IRMutatorWithAnalyzer::VisitStmt_(op));
939
940
    if (result_.for_map.count(tvm::ffi::GetRef<For>(op))) {
      auto root = tvm::ffi::GetRef<For>(op);
941
942
943
944
945
946
      // 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.
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
      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
964
965
      // We use PostOrderVisit to detect whether the loop only manuplates
      // "local" buffers, which indicates register usage and justifies skipping
966
      // thread binding.
967
      bool local_register_only = true;
968
969
      PostOrderVisit(root, [&](const ObjectRef &obj) {
        if (const auto *store = obj.as<BufferStoreNode>()) {
970
971
972
973
974
975
          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;
976
977
978
979
          }
        }
      });

980
      auto loop_layout = result_.for_map[root];
981
      // FIXME: tell in-Parallel and out-of-Parallel `local`s apart
982
983
984
      // 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;
985

986
      if (parallel_loop) {
987
988
989
        for_node =
            PartitionLoop(for_node, thread_var_->var, analyzer_, loop_layout);
      }
990
      // If none thread bindings are provided, partition the loop
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
      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;
          }
        }
      });
1005
1006
1007
1008
1009
1010
1011
1012
1013
      // 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;
          }
      });
1014

1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
      // 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) {
1027
1028
        for_node = VectorizeLoop(for_node);
      }
1029

1030
1031
      if (result_.predicate_map.count(root) && parallel_loop) {
        return IfThenElse(result_.predicate_map[root], for_node);
1032
1033
1034
1035
1036
1037
1038
      } else {
        return for_node;
      }
    }
    return for_node;
  }

1039
  Stmt VisitStmt_(const AttrStmtNode *op) final {
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
    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);
  }

1050
private:
1051
  const LayoutInferenceResult result_;
1052
1053
  IterVar thread_var_ = IterVar(Range::FromMinExtent(0, 1), Var("v_thread"),
                                IterVarType::kDataPar);
1054
  bool skip_thread_partition_{false};
1055
1056
1057
1058
};

tvm::transform::Pass LayoutInference() {
  using namespace tir::transform;
1059
  auto pass_func = [=](PrimFunc f, const IRModule &m, const PassContext &ctx) {
1060
    f.CopyOnWrite()->body = ParallelLoopTransformer::Substitute(f->body);
1061
    ThreadBindingCollector collector;
1062
    collector(f->body);
1063
    bool has_thread_binding = !collector.thread_binding_.empty();
1064
    bool skip_thread_partition = !has_thread_binding;
1065
    return LayoutInferencer::Substitute(std::move(f), skip_thread_partition);
1066
1067
1068
1069
  };
  return CreatePrimFuncPass(pass_func, 0, "tl.LayoutInference", {});
}

1070
TVM_FFI_STATIC_INIT_BLOCK() {
1071
1072
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def("tl.transform.LayoutInference", LayoutInference);
1073
}
1074

1075
1076
} // namespace tl
} // namespace tvm