lower_tile_op.cc 22.1 KB
Newer Older
1
2
3
4
5
/*!
 * \file lower_tile_op.cc
 * \brief Lower the tile op for further codegen.
 */

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

#include "../layout/layout.h"
#include "../layout/utils.h"
14
#include "../op/builtin.h"
15
16
#include "../op/gemm.h"
#include "../op/gemm_sp.h"
17
#include "../op/operator.h"
18

19
#include "arith/ir_mutator_with_analyzer.h"
20
21
22
23
24
25
26
#include "loop_partition.h"

namespace tvm {
namespace tl {

using namespace tir;

27
28
static Buffer makeBufferWithLayout(const Buffer &buffer, const Layout &layout,
                                   Map<Var, Var> &var_remap) {
29
30
  const auto *ptr_type =
      TVM_TYPE_AS(buffer->data->type_annotation, PointerTypeNode);
31
32
33
34
35
36
37
38
39
40
41
  Type new_type;
  // convert fragments to normal local buffer
  if (ptr_type->storage_scope == "local.fragment") {
    new_type = PointerType(ptr_type->element_type, "local");
  } else {
    new_type = buffer->data->type_annotation;
  }
  Var new_var;
  if (ptr_type->storage_scope == "global") {
    new_var = buffer->data;
  } else {
42
43
44
45
46
47
    if (var_remap.count(buffer->data)) {
      new_var = var_remap[buffer->data];
    } else {
      new_var = Var(buffer->data->name_hint, new_type);
      var_remap.Set(buffer->data, new_var);
    }
48
  }
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
  Array<PrimExpr> layout_shape = layout->OutputShape();
  Array<PrimExpr> output_shape = layout_shape;

  if (ptr_type->storage_scope == "shared" ||
      ptr_type->storage_scope == "shared.dyn") {
    int replicate_extent = 1;
    Array<PrimExpr> buffer_shape = buffer->shape;
    int buffer_extent = 1;
    int layout_extent = 1;
    for (size_t i = 0; i < buffer_shape.size(); i++) {
      auto shape = buffer_shape[i].as<IntImmNode>();
      buffer_extent *= shape->value;
    }
    for (size_t i = 0; i < layout_shape.size(); i++) {
      auto shape = layout_shape[i].as<IntImmNode>();
      layout_extent *= shape->value;
    }
    replicate_extent = buffer_extent / layout_extent;
    if (replicate_extent > 1) {
      output_shape.insert(output_shape.begin(), replicate_extent);
    }
  }
  return Buffer(new_var, buffer->dtype, output_shape, {}, buffer->elem_offset,
                buffer->name, buffer->data_alignment, buffer->offset_factor,
                buffer->buffer_type);
74
75
}

76
77
78
79
80
81
class BufferGemmCollector : public StmtExprVisitor {
public:
  BufferGemmCollector() { Clear(); }

  void Clear() { buffer_var_gemm_.clear(); }

82
  void Collect(const Stmt &stmt) { VisitStmt(stmt); }
83
84
85
86
87
88

  Array<Var> GetBufferVarGemm() { return buffer_var_gemm_; }

private:
  void VisitStmt_(const EvaluateNode *op) {
    auto call = Downcast<Call>(op->value);
89
    if (call->op.same_as(Gemm::Get())) {
90
91
92
93
94
95
96
97
98
99
100
101
      auto srcA_buffer_access_ptr = Downcast<Call>(call->args[0]);
      ICHECK(srcA_buffer_access_ptr->op.same_as(builtin::tvm_access_ptr()));
      auto srcA_buffer_var = Downcast<Var>(srcA_buffer_access_ptr->args[1]);
      auto srcB_buffer_access_ptr = Downcast<Call>(call->args[1]);
      ICHECK(srcB_buffer_access_ptr->op.same_as(builtin::tvm_access_ptr()));
      auto srcB_buffer_var = Downcast<Var>(srcB_buffer_access_ptr->args[1]);
      auto dst_buffer_access_ptr = Downcast<Call>(call->args[2]);
      ICHECK(dst_buffer_access_ptr->op.same_as(builtin::tvm_access_ptr()));
      auto dst_buffer_var = Downcast<Var>(dst_buffer_access_ptr->args[1]);
      buffer_var_gemm_.push_back(srcA_buffer_var);
      buffer_var_gemm_.push_back(srcB_buffer_var);
      buffer_var_gemm_.push_back(dst_buffer_var);
102
    } else if (call->op.same_as(GemmSP::Get())) {
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
      auto srcA_buffer_access_ptr = Downcast<Call>(call->args[0]);
      ICHECK(srcA_buffer_access_ptr->op.same_as(builtin::tvm_access_ptr()));
      auto srcA_buffer_var = Downcast<Var>(srcA_buffer_access_ptr->args[1]);
      auto srcB_buffer_access_ptr = Downcast<Call>(call->args[1]);
      ICHECK(srcB_buffer_access_ptr->op.same_as(builtin::tvm_access_ptr()));
      auto srcB_buffer_var = Downcast<Var>(srcB_buffer_access_ptr->args[1]);
      auto dst_buffer_access_ptr = Downcast<Call>(call->args[2]);
      ICHECK(dst_buffer_access_ptr->op.same_as(builtin::tvm_access_ptr()));
      auto dst_buffer_var = Downcast<Var>(dst_buffer_access_ptr->args[1]);
      buffer_var_gemm_.push_back(srcA_buffer_var);
      buffer_var_gemm_.push_back(srcB_buffer_var);
      buffer_var_gemm_.push_back(dst_buffer_var);
    }
  }

  Array<Var> buffer_var_gemm_;
};

121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
/*!
 * \brief A class that rewrites buffer references in a statement based on a
 * given buffer remapping.
 *
 * This class is used to update buffer references in a statement after buffer
 * transformations have been applied. It specifically handles the remapping of
 * padding annotations.
 */
class RemapBufferRewriter : public arith::IRMutatorWithAnalyzer {
public:
  /*!
   * \brief Substitute buffer references in a statement based on a given buffer
   * remapping. \param stmt The statement to rewrite. \param buffer_remap A map
   * from old buffers to new buffers. \return The rewritten statement.
   */
136
  static Stmt Substitute(const Stmt &stmt, Map<Buffer, Buffer> buffer_remap) {
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
    arith::Analyzer analyzer;
    RemapBufferRewriter substituter(&analyzer);
    substituter.buffer_remap_ = std::move(buffer_remap);
    return substituter.VisitStmt(stmt);
  }

private:
  using arith::IRMutatorWithAnalyzer::IRMutatorWithAnalyzer;

  Stmt VisitStmt_(const BlockNode *op) final {
    if (op->annotations.count(attr::kPaddingMap)) {
      return RewritePaddingMap(op);
    }
    return IRMutatorWithAnalyzer::VisitStmt_(op);
  }

  /*!
   * \brief Rewrite the padding map annotation of a block.
   * \param op The block node to rewrite.
   * \return The rewritten block.
   */
  Stmt RewritePaddingMap(const BlockNode *op) {
159
160
161
162
    auto padding_map = op->annotations.Get(attr::kPaddingMap);
    if (!padding_map) {
      LOG(FATAL) << "Padding map annotation is missing";
    }
163
164

    Map<Var, Var> var_remap = CreateVarRemap();
165
166
    Map<Var, PrimExpr> new_padding_map = RemapPaddingMap(
        Downcast<Map<Var, PrimExpr>>(padding_map.value()), var_remap);
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207

    auto block = Downcast<Block>(IRMutatorWithAnalyzer::VisitStmt_(op));
    auto block_ptr = block.CopyOnWrite();
    block_ptr->annotations.Set(attr::kPaddingMap, new_padding_map);
    return block;
  }

  /*!
   * \brief Create a mapping from old variables to new variables based on buffer
   * remapping. \return A map from old variables to new variables.
   */
  Map<Var, Var> CreateVarRemap() const {
    Map<Var, Var> var_remap;
    for (const auto &[buffer, buffer_remap] : buffer_remap_) {
      var_remap.Set(buffer->data, buffer_remap->data);
    }
    return var_remap;
  }

  /*!
   * \brief Remap the padding map using the variable remapping.
   * \param padding_map The original padding map.
   * \param var_remap The variable remapping.
   * \return The remapped padding map.
   */
  Map<Var, PrimExpr> RemapPaddingMap(const Map<Var, PrimExpr> &padding_map,
                                     const Map<Var, Var> &var_remap) const {
    Map<Var, PrimExpr> new_padding_map;
    for (const auto &[var, padding] : padding_map) {
      if (var_remap.count(var)) {
        new_padding_map.Set(var_remap.at(var), padding);
      } else {
        new_padding_map.Set(var, padding);
      }
    }
    return new_padding_map;
  }

  Map<Buffer, Buffer> buffer_remap_;
};

208
class LowerTileOpPass : arith::IRMutatorWithAnalyzer {
209
public:
210
211
212
213
214
  static PrimFunc Substitute(PrimFunc f) {
    arith::Analyzer analyzer;
    LowerTileOpPass substituter(&analyzer);
    // Trace the buffer map for tvm_access_ptr
    substituter.buffer_map_.insert(f->buffer_map.begin(), f->buffer_map.end());
215
    for (const auto &[_, buffer] : f->buffer_map) {
216
217
218
219
220
      substituter.buffer_data_to_buffer_.Set(buffer->data, buffer);
    }
    auto target = f->GetAttr<Target>(tvm::attr::kTarget);
    ICHECK(target.defined()) << "LowerTileOpPass: Require the target attribute";
    substituter.target_ = target.value();
221
222
223
224
225
    // For TMA 1D, we should collect the buffers which are not used in GEMM and
    // do not need swizzle
    BufferGemmCollector collector;
    collector.Collect(f->body);
    substituter.buffer_var_gemm_ = collector.GetBufferVarGemm();
226
    PrimFuncNode *fptr = f.CopyOnWrite();
227
    fptr->body = substituter.VisitStmt(f->body);
228
229
    fptr->body =
        RemapBufferRewriter::Substitute(fptr->body, substituter.buffer_remap_);
230
231
232
233
234
235
236
237
238
    tvm::transform::PassContext ctxt = tvm::transform::PassContext::Current();
    Optional<Bool> opt_disable_tma_lower =
        ctxt->GetConfig(kDisableTMALower, Optional<Bool>());

    if (!opt_disable_tma_lower.value_or(Bool(false))) {
      // @lei: this is a workaround, as if we don't disable tma lower,
      // cp async lowering won't be generated.
      ctxt->config.Set(kDisableTMALower, Bool(!substituter.has_tma_));
    }
239
240
241
    return f;
  }

242
private:
243
244
  using arith::IRMutatorWithAnalyzer::IRMutatorWithAnalyzer;

245
  Stmt VisitStmt_(const BlockNode *op) final {
246
247
248
249
250
251
252
253
254
255
256
257
    // Record the mapping from buffer data var to buffer for later lookup
    for (auto buffer : op->alloc_buffers) {
      buffer_map_.insert({buffer->data, buffer});
    }
    for (auto match_buffer : op->match_buffers) {
      buffer_map_.insert({match_buffer->buffer->data, match_buffer->buffer});
    }
    for (auto buffer : op->alloc_buffers) {
      buffer_data_to_buffer_.Set(buffer->data, buffer);
    }
    Map<Var, Layout> vmap;
    if (op->annotations.count(attr::kLayoutMap)) {
258
259
260
      auto layout_map = op->annotations.at(attr::kLayoutMap)
                            .as<Map<Buffer, Layout>>()
                            .value();
261
      for (auto [buffer, layout] : layout_map) {
262
263
        buffer_remap_.Set(buffer,
                          makeBufferWithLayout(buffer, layout, var_remap_));
264
265
266
267
268
269
270
271
272
273
274
        layout_map_.Set(buffer, layout);
      }
    }
    auto block = Downcast<Block>(arith::IRMutatorWithAnalyzer::VisitStmt_(op));
    auto block_ptr = block.CopyOnWrite();
    for (size_t i = 0; i < block->alloc_buffers.size(); i++) {
      auto buffer = block->alloc_buffers[i];
      if (buffer_remap_.count(buffer)) {
        block_ptr->alloc_buffers.Set(i, buffer_remap_[buffer]);
      }
    }
275
276
    for (const auto &buffer : workspaces_)
      block_ptr->alloc_buffers.push_back(buffer);
277
278
279
280
281
    workspaces_.clear();
    block_ptr->annotations.erase(attr::kLayoutMap);
    return block;
  }

282
  int CheckAndGetBufferRowSize(const Buffer &buffer) {
283
    CHECK(buffer->shape.size() >= 2)
284
285
        << "The dimension of Buffer \"" << buffer->name << "\" with shape "
        << buffer->shape << " should be at least 2";
286
287
288
289
290
291

    auto dim = buffer->shape.size();
    auto buffer_row_size = buffer->shape[dim - 1].as<IntImmNode>()->value;
    return buffer_row_size;
  }

292
293
294
295
  PrimExpr
  HandleAccessPtrAndOffset(const PrimExpr &access_ptr,
                           const Optional<PrimExpr> &offset = std::nullopt,
                           DataType dtype = DataType::Int(32)) {
296
297
    // The 2th arg of T.tvm_access_ptr call is offset, we set it to 0 and
    // accumulate it to smem_offset
298
299
300
301
302
303
304
305
306
307
308
    CHECK(access_ptr->IsInstance<CallNode>())
        << "Invalid access ptr for permuted layout: " << access_ptr;
    auto access_ptr_call = Downcast<Call>(access_ptr);
    if (access_ptr_call->op.same_as(builtin::tvm_access_ptr())) {
      LOG(FATAL) << "Transformation for tvm_access_ptr is not implemented yet";
    } else if (access_ptr_call->op.same_as(builtin::address_of())) {
      BufferLoad load = Downcast<BufferLoad>(access_ptr_call->args[0]);
      Array<PrimExpr> indices = load->indices;
      Array<PrimExpr> shape = load->buffer->shape;

      CHECK_EQ(indices.size(), shape.size())
309
310
311
312
          << "Indices size and shape size must match for general N-dimensional "
             "buffer "
          << "but got indices size: " << indices.size()
          << " and shape size: " << shape.size();
313
314
315
316
317
318
319
320
321

      PrimExpr elem_offset = 0;
      PrimExpr stride = 1;

      for (int i = static_cast<int>(shape.size()) - 1; i >= 0; --i) {
        elem_offset += indices[i] * stride;
        stride *= shape[i];
      }

322
323
      PrimExpr smem_offset =
          elem_offset + (offset.defined() ? offset.value() : 0);
324
325
326

      auto new_buffer = buffer_remap_[load->buffer];

327
328
      auto buffer_map_iter =
          buffer_map_.find(Downcast<Var>(load->buffer->data));
329
      CHECK(buffer_map_iter != buffer_map_.end())
330
331
          << "The buffer corresponding to data Var " << access_ptr_call->args[0]
          << " is not found";
332
333
334
335
336
337
338
339
340

      int buffer_row_size = CheckAndGetBufferRowSize(buffer_map_iter->second);
      (void)buffer_row_size;

      // Convert offset to target-dimension, reindex it and convert it back
      Array<PrimExpr> multi_dim_indices;
      PrimExpr remaining_offset = smem_offset;

      for (int i = static_cast<int>(shape.size()) - 1; i >= 0; --i) {
341
342
        multi_dim_indices.insert(multi_dim_indices.begin(),
                                 floormod(remaining_offset, shape[i]));
343
344
345
        remaining_offset = floordiv(remaining_offset, shape[i]);
      }

346
347
      auto forward_indices =
          layout_map_[load->buffer]->Forward(multi_dim_indices);
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
      PrimExpr new_offset = 0;
      PrimExpr stride_offset = 1;
      for (int i = static_cast<int>(shape.size()) - 1; i >= 0; --i) {
        new_offset += forward_indices[i] * stride_offset;
        stride_offset *= shape[i];
      }
      new_offset = analyzer_->Simplify(new_offset);

      Array<PrimExpr> new_indices;
      for (int i = static_cast<int>(shape.size()) - 1; i >= 0; --i) {
        new_indices.insert(new_indices.begin(), floormod(new_offset, shape[i]));
        new_offset = floordiv(new_offset, shape[i]);
      }

      auto new_access_ptr = access_ptr_call.CopyOnWrite();
      new_access_ptr->args.Set(0, BufferLoad(new_buffer, new_indices));
    } else {
      LOG(FATAL) << "Invalid access op for permuted layout: " << access_ptr;
    }

    return access_ptr_call;
  }

371
  PrimExpr VisitExpr_(const tir::CallNode *op) final {
372
373
374
375
376
    if ((!has_tma_) && (op->op.same_as(tl::tma_load()) ||
                        op->op.same_as(tl::tma_load_im2col()) ||
                        op->op.same_as(tl::tma_store()))) {
      has_tma_ = true;
    }
377
    Array<RelaxExpr> ptx_instructions = {builtin::ptx_ldmatrix(),
378
379
380
381
382
383
                                         builtin::mma_store()};

    if (std::find(ptx_instructions.begin(), ptx_instructions.end(), op->op) ==
        ptx_instructions.end()) {
      auto call = Downcast<Call>(IRMutatorWithAnalyzer::VisitExpr_(op));
      return call;
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
    } else {
      is_ptx_ = true;
    }
    // Rewrite from/to shared or shared.dyn to/from local
    auto call = Downcast<Call>(IRMutatorWithAnalyzer::VisitExpr_(op));
    if (call->op.same_as(builtin::ptx_ldmatrix())) {
      // form: T.ptx_ldmatrix(..., smem_ptr, smem_offset)
      // smem_ptr: T.tvm_access_ptr(ptype, data, offset, extent, rw_mask)
      // or T.address_of(buffer, offset)
      auto access_ptr = call->args[5];
      PrimExpr smem_offset = call->args[6];
      Call address_of_call = Downcast<Call>(access_ptr);
      if (!address_of_call->op.same_as(builtin::address_of())) {
        LOG(FATAL) << "Invalid access ptr for permuted layout: " << access_ptr;
      }
      BufferLoad load = Downcast<BufferLoad>(address_of_call->args[0]);

      if (buffer_remap_.count(load->buffer)) {
402
403
        auto new_access_ptr =
            HandleAccessPtrAndOffset(access_ptr, smem_offset, call->dtype);
404
405
406
407
408
        auto new_call = call.CopyOnWrite();
        new_call->args.Set(5, new_access_ptr);
        new_call->args.Set(6, IntImm(smem_offset->dtype, 0));
      }
    } else if (call->op.same_as(builtin::mma_store())) {
409
410
      // because we will directly store result to Buffer instead of calling
      // mma_store now
411
      auto access_ptr = call->args[2];
412
      auto new_access_ptr =
413
          HandleAccessPtrAndOffset(access_ptr, std::nullopt, call->dtype);
414
415
416
417
418
419
420
421
422
      auto new_call = call.CopyOnWrite();
      new_call->args.Set(2, new_access_ptr);
    } else {
      LOG(FATAL) << "Invalid call node: " << call;
    }
    is_ptx_ = false;
    return call;
  }

423
  PrimExpr VisitExpr_(const BufferLoadNode *op) final {
424
425
426
427
    auto load = Downcast<BufferLoad>(IRMutatorWithAnalyzer::VisitExpr_(op));
    if (is_ptx_) {
      return load;
    }
428
429
430
    auto buffer = load->buffer;
    if (buffer_remap_.count(buffer)) {
      auto new_indices = layout_map_[buffer]->Forward(load->indices);
431
432
      auto new_buffer = buffer_remap_[load->buffer];
      return BufferLoad(new_buffer, new_indices);
433
434
435
436
437
438
    } else if (var_remap_.count(buffer->data)) {
      auto new_buffer = Buffer(
          var_remap_[buffer->data], buffer->dtype, buffer->shape,
          buffer->strides, buffer->elem_offset, buffer->name,
          buffer->data_alignment, buffer->offset_factor, buffer->buffer_type);
      return BufferLoad(new_buffer, load->indices);
439
440
441
442
    }
    return load;
  }

443
  Stmt VisitStmt_(const BufferStoreNode *op) final {
444
    auto store = Downcast<BufferStore>(IRMutatorWithAnalyzer::VisitStmt_(op));
445
446
447
    auto buffer = store->buffer;
    if (buffer_remap_.count(buffer)) {
      auto new_indices = layout_map_[buffer]->Forward(store->indices);
448
449
      auto new_buffer = buffer_remap_[store->buffer];
      return BufferStore(new_buffer, store->value, new_indices);
450
451
452
453
454
455
    } else if (var_remap_.count(buffer->data)) {
      auto new_buffer = Buffer(
          var_remap_[buffer->data], buffer->dtype, buffer->shape,
          buffer->strides, buffer->elem_offset, buffer->name,
          buffer->data_alignment, buffer->offset_factor, buffer->buffer_type);
      return BufferStore(new_buffer, store->value, store->indices);
456
457
458
459
    }
    return store;
  }

460
  PrimExpr VisitExpr_(const VarNode *op) final {
461
462
463
    auto var = Downcast<Var>(IRMutatorWithAnalyzer::VisitExpr_(op));
    if (buffer_data_to_buffer_.count(var)) {
      auto buffer = buffer_data_to_buffer_[var];
464
465
      if (buffer_remap_.count(buffer))
        return buffer_remap_[buffer]->data;
466
467
468
469
    }
    return var;
  }

470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
  /**
   * @brief Handle an Evaluate node, lowering a detected tile operator to TIR.
   *
   * This visit implementation detects whether the Evaluate node represents a
   * tile operator invocation (via ParseOperator). If no tile operator is found
   * or the call targets a global function, the node is delegated to the base
   * visitor.
   *
   * When a tile operator is present, the method:
   * - Builds a workspace-allocation callback that creates a dynamic shared
   * buffer named "workspace" (storage scope "shared.dyn") and returns its write
   *   access pointer.
   * - Determines thread bounds for lowering from the analyzer's constant-int
   *   information for thread_var_; if unavailable, a default range [0,1) is
   * used.
   * - Invokes tile_op->Lower(...) with LowerArgs containing target, thread
   *   bounds, thread variable, the workspace callback, layout and buffer remap
   *   maps, and the list of GEMM-involved buffer vars; the analyzer is passed
   *   through for use during lowering.
   *
   * The lowered statement returned by the operator is then visited by the base
   * IRMutatorWithAnalyzer and that result is returned.
   *
   * @return Stmt The (possibly transformed) statement after lowering or base
   * visitor processing.
   */
496
  Stmt VisitStmt_(const EvaluateNode *op) final {
497
    // LOG(INFO) << "evaluate node: " << op->value;
498
    const CallNode *call = op->value.as<CallNode>();
499
    // LOG(INFO) << "call: " << call->op;
500
501
502
503
504
    // Do not analysis the call node to the global function.
    if (call && call->op.as<GlobalVarNode>())
      return Downcast<Evaluate>(IRMutatorWithAnalyzer::VisitStmt_(op));

    auto tile_op = ParseOperator(GetRef<Stmt>(op), buffer_data_to_buffer_);
505
    if (!tile_op.defined())
506
      return IRMutatorWithAnalyzer::VisitStmt_(op);
507
    AddWorkspaceCallback callback = [this](int num_elem, DataType dtype) {
508
509
      auto workspace =
          decl_buffer({PrimExpr(num_elem)}, dtype, "workspace", "shared.dyn");
510
      workspaces_.push_back(workspace);
511
      return workspace.access_ptr(2); // write
512
513
    };

514
515
516
517
518
519
    Range thread_bounds;

    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;
520
      auto extent = max_value + 1 - min_value;
521
522
      thread_bounds =
          Range::FromMinExtent(IntImm(thread_var_->var.dtype(), min_value),
523
                               IntImm(thread_var_->var.dtype(), extent));
524
525
526
    } else {
      thread_bounds = Range::FromMinExtent(0, 1);
    }
527

528
529
530
531
    auto lowered = tile_op->Lower(
        LowerArgs{target_, thread_bounds, thread_var_->var, callback,
                  layout_map_, buffer_remap_, buffer_var_gemm_},
        analyzer_);
532
533
534
    return IRMutatorWithAnalyzer::VisitStmt(lowered);
  }

535
  Stmt VisitStmt_(const AttrStmtNode *op) final {
536
537
538
539
    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") {
540
        thread_var_ = iv;
541
542
543
544
545
546
547
548
549
550
551
        ICHECK(iv->dom->extent.as<IntImmNode>());
        thread_block_size_ = iv->dom->extent.as<IntImmNode>()->value;
      }
    }
    return arith::IRMutatorWithAnalyzer::VisitStmt_(op);
  }

  Target target_;
  Map<Var, Buffer> buffer_data_to_buffer_;
  Map<Buffer, Layout> layout_map_;
  Map<Buffer, Buffer> buffer_remap_;
552
553
554
555
  // 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);
556
557
558
559
560
561
562
  size_t thread_block_size_ = 0;
  Array<Buffer> workspaces_;
  // For ptx Node, we need to remap the buffer and indices
  // By access CallNode instead of BufferLoad Node.
  bool is_ptx_{false};
  // Mapping from data Var of a Buffer to Buffer, for lookup
  std::unordered_map<Var, Buffer, ObjectPtrHash, ObjectPtrEqual> buffer_map_;
563
  Map<Var, Var> var_remap_;
564
  bool has_tma_{false};
565
  Array<Var> buffer_var_gemm_;
566
567
568
569
570
571
572
};

namespace transform {

using namespace tir::transform;

tvm::transform::Pass LowerTileOp() {
573
  auto pass_func = [=](PrimFunc f, const IRModule &m, const PassContext &ctx) {
574
575
576
577
578
    return LowerTileOpPass::Substitute(std::move(f));
  };
  return CreatePrimFuncPass(pass_func, 0, "tl.LowerTileOp", {});
}

579
580
581
582
TVM_FFI_STATIC_INIT_BLOCK({
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def("tl.transform.LowerTileOp", LowerTileOp);
});
583
} // namespace transform
584

585
586
} // namespace tl
} // namespace tvm