"docs/vscode:/vscode.git/clone" did not exist on "5529d58b1b96de19fb1088ecd46426914cc1e926"
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
    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;
306
      Array<PrimExpr> old_shape = load->buffer->shape;
307

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

      PrimExpr elem_offset = 0;
      PrimExpr stride = 1;

317
      for (int i = static_cast<int>(old_shape.size()) - 1; i >= 0; --i) {
318
        elem_offset += indices[i] * stride;
319
        stride *= old_shape[i];
320
321
      }

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

      auto new_buffer = buffer_remap_[load->buffer];
326
      auto new_shape = new_buffer->shape;
327

328
329
      auto buffer_map_iter =
          buffer_map_.find(Downcast<Var>(load->buffer->data));
330
      CHECK(buffer_map_iter != buffer_map_.end())
331
332
          << "The buffer corresponding to data Var " << access_ptr_call->args[0]
          << " is not found";
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;

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

347
348
      auto forward_indices =
          layout_map_[load->buffer]->Forward(multi_dim_indices);
349
350
      PrimExpr new_offset = 0;
      PrimExpr stride_offset = 1;
351
      for (int i = static_cast<int>(new_shape.size()) - 1; i >= 0; --i) {
352
        new_offset += forward_indices[i] * stride_offset;
353
        stride_offset *= new_shape[i];
354
355
356
357
      }
      new_offset = analyzer_->Simplify(new_offset);

      Array<PrimExpr> new_indices;
358
359
360
361
      for (int i = static_cast<int>(new_shape.size()) - 1; i >= 0; --i) {
        new_indices.insert(new_indices.begin(),
                           floormod(new_offset, new_shape[i]));
        new_offset = floordiv(new_offset, new_shape[i]);
362
363
364
365
366
367
368
369
370
371
372
      }

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

373
  PrimExpr VisitExpr_(const tir::CallNode *op) final {
374
375
376
377
378
    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;
    }
379
    Array<RelaxExpr> ptx_instructions = {builtin::ptx_ldmatrix(),
380
381
382
383
384
385
                                         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;
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
    } 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)) {
403
404
        auto new_access_ptr =
            HandleAccessPtrAndOffset(access_ptr, smem_offset, call->dtype);
405
406
407
408
409
        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())) {
410
411
      // because we will directly store result to Buffer instead of calling
      // mma_store now
412
      auto access_ptr = call->args[2];
413
      auto new_access_ptr =
414
          HandleAccessPtrAndOffset(access_ptr, std::nullopt, call->dtype);
415
416
417
418
419
420
421
422
423
      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;
  }

424
  PrimExpr VisitExpr_(const BufferLoadNode *op) final {
425
426
427
428
    auto load = Downcast<BufferLoad>(IRMutatorWithAnalyzer::VisitExpr_(op));
    if (is_ptx_) {
      return load;
    }
429
430
431
    auto buffer = load->buffer;
    if (buffer_remap_.count(buffer)) {
      auto new_indices = layout_map_[buffer]->Forward(load->indices);
432
433
      auto new_buffer = buffer_remap_[load->buffer];
      return BufferLoad(new_buffer, new_indices);
434
435
436
437
438
439
    } 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);
440
441
442
443
    }
    return load;
  }

444
  Stmt VisitStmt_(const BufferStoreNode *op) final {
445
    auto store = Downcast<BufferStore>(IRMutatorWithAnalyzer::VisitStmt_(op));
446
447
448
    auto buffer = store->buffer;
    if (buffer_remap_.count(buffer)) {
      auto new_indices = layout_map_[buffer]->Forward(store->indices);
449
450
      auto new_buffer = buffer_remap_[store->buffer];
      return BufferStore(new_buffer, store->value, new_indices);
451
452
453
454
455
456
    } 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);
457
458
459
460
    }
    return store;
  }

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

471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
  /**
   * @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.
   */
497
498
  Stmt VisitStmt_(const EvaluateNode *op) final {
    const CallNode *call = op->value.as<CallNode>();
499
500
501
502
503
    // 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_);
504
    if (!tile_op.defined())
505
      return IRMutatorWithAnalyzer::VisitStmt_(op);
506
    AddWorkspaceCallback callback = [this](int num_elem, DataType dtype) {
507
508
      auto workspace =
          decl_buffer({PrimExpr(num_elem)}, dtype, "workspace", "shared.dyn");
509
      workspaces_.push_back(workspace);
510
      return workspace.access_ptr(2); // write
511
512
    };

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

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

534
  Stmt VisitStmt_(const AttrStmtNode *op) final {
535
536
537
538
    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") {
539
        thread_var_ = iv;
540
541
542
543
544
545
546
547
548
549
550
        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_;
551
552
553
554
  // 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);
555
556
557
558
559
560
561
  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_;
562
  Map<Var, Var> var_remap_;
563
  bool has_tma_{false};
564
  Array<Var> buffer_var_gemm_;
565
566
567
568
569
570
571
};

namespace transform {

using namespace tir::transform;

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

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

584
585
} // namespace tl
} // namespace tvm