gemm_layouts.cc 19.9 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
/*!
 * \file layout/gemm_layouts.cc
 * \brief Define Layout used in MMA and other operations.
 *
 */

#include <tvm/tir/stmt_functor.h>

#include <cmath>

#include "layout.h"

namespace tvm {
namespace tl {

static IterVar make_itervar(std::string name, PrimExpr dom) {
17
  Var var = Var(name, dom->dtype);
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
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
74
75
76
77
78
  return IterVar(Range(0, dom), var, IterVarType::kDataPar);
}

Fragment makeGemmFragment8x8() {
  IterVar i = make_itervar("i", 8);
  IterVar j = make_itervar("j", 8);
  IterVar rep = make_itervar("rep", 1);
  PrimExpr forward_thread = FloorDiv(j->var, 2) + 4 * i;
  PrimExpr index = FloorMod(j->var, 2);
  return Fragment({i, j}, {index}, forward_thread, rep);
}
/*
From https://github.com/RadeonOpenCompute/amd_matrix_instruction_calculator
./matrix_calculator.py --architecture cdna1 --instruction v_mfma_f32_16x16x16f16
--detail-instruction
*/
Fragment makeGemmFragmentAB16x16CDNA() {
  IterVar i = make_itervar("i", 16);
  IterVar j = make_itervar("j", 16);
  IterVar rep = make_itervar("rep", 1);
  PrimExpr forward_thread = 16 * FloorDiv(j->var, 4) + i;
  PrimExpr index = FloorMod(j->var, 4);
  return Fragment({i, j}, {index}, forward_thread, rep);
}

Fragment makeGemmFragmentAB16x16CDNATransposed() {
  IterVar i = make_itervar("i", 16);
  IterVar j = make_itervar("j", 16);
  IterVar rep = make_itervar("rep", 1);
  PrimExpr forward_thread = 16 * FloorDiv(i->var, 4) + j;
  PrimExpr index = FloorMod(i->var, 4);
  return Fragment({i, j}, {index}, forward_thread, rep);
}

Fragment makeGemmFragmentC16x16CDNA() {
  IterVar i = make_itervar("i", 16);
  IterVar j = make_itervar("j", 16);
  IterVar rep = make_itervar("rep", 1);
  PrimExpr forward_thread = 16 * FloorDiv(j->var, 4) + i;
  PrimExpr index = FloorMod(j->var, 4);
  return Fragment({i, j}, {index}, forward_thread, rep);
}

Fragment makeGemmFragment8x8Transposed() {
  IterVar i = make_itervar("i", 8);
  IterVar j = make_itervar("j", 8);
  IterVar rep = make_itervar("rep", 1);
  PrimExpr forward_thread = FloorDiv(i->var, 2) + 4 * j;
  PrimExpr index = FloorMod(i->var, 2);
  return Fragment({i, j}, {index}, forward_thread, rep);
}

Fragment makeGemmFragment8x16() {
  IterVar i = make_itervar("i", 8);
  IterVar j = make_itervar("j", 16);
  IterVar rep = make_itervar("rep", 1);
  PrimExpr forward_thread = FloorDiv(j->var, 4) + 4 * i;
  PrimExpr index = FloorMod(j->var, 4);
  return Fragment({i, j}, {index}, forward_thread, rep);
}

79
80
Fragment makeGemmFragmentC_F64(const int block_m, const int block_n,
                               const int warp_m, const int warp_n) {
81
82
83
  ICHECK(block_m % warp_m == 0);
  ICHECK(block_n % warp_n == 0);
  ICHECK(warp_m % 16 == 0);
84
  ICHECK(warp_n % 8 == 0);
85
  auto base_layout = makeGemmFragment8x8();
86
87
88
89
  auto warp_layout =
      base_layout->Repeat({block_m / warp_m, block_n / warp_n}, true, false);
  auto block_layout =
      warp_layout->Repeat({warp_m / 8, warp_n / 8}, false, false);
90
91
92
  return block_layout;
}

93
94
Fragment makeGemmFragmentC(const int block_m, const int block_n,
                           const int warp_m, const int warp_n,
95
                           const int element_size) {
96
97
  if (element_size == 64)
    return makeGemmFragmentC_F64(block_m, block_n, warp_m, warp_n);
98
99
100
  ICHECK(block_m % warp_m == 0);
  ICHECK(block_n % warp_n == 0);
  ICHECK(warp_m % 16 == 0) << "warp_m=" << warp_m;
101
  ICHECK(warp_n % 8 == 0) << "warp_n=" << warp_n;
102
  auto base_layout = makeGemmFragment8x8()->Repeat({2, 1}, false);
103
104
105
106
  auto warp_layout =
      base_layout->Repeat({block_m / warp_m, block_n / warp_n}, true, false);
  auto block_layout =
      warp_layout->Repeat({warp_m / 16, warp_n / 8}, false, false);
107
108
109
  return block_layout;
}

110
111
112
113
114
Fragment makeGemmFragmentCCDNA(const int block_m, const int block_n,
                               const int warp_m, const int warp_n,
                               const int element_size) {
  if (element_size == 64)
    LOG(FATAL) << "Not supported";
115
116
117
118
119
  ICHECK(block_m % warp_m == 0);
  ICHECK(block_n % warp_n == 0);
  ICHECK(warp_m % 16 == 0) << "warp_m=" << warp_m;
  ICHECK(warp_n % 16 == 0) << "warp_n=" << warp_n;
  auto base_layout = makeGemmFragmentC16x16CDNA()->Repeat({1, 1}, false);
120
121
122
123
  auto warp_layout =
      base_layout->Repeat({warp_m / 16, warp_n / 16}, false, true);
  auto block_layout =
      warp_layout->Repeat({block_m / warp_m, block_n / warp_n}, true, false);
124
125
126
  return block_layout;
}

127
128
129
Fragment makeGemmFragmentCHopper(const int block_m, const int block_n,
                                 const int warp_m, const int warp_n,
                                 const int element_size) {
130
131
  ICHECK(block_m % warp_m == 0);
  // ICHECK(block_n == warp_n);
132
  ICHECK(warp_m % 16 == 0) << "warp_m=" << warp_m;
133
134
135
136
  auto warp_layout = makeGemmFragment8x8()->Repeat({2, warp_n / 8}, false,
                                                   false); // 16 x N (1 warp)
  auto block_layout = warp_layout->Repeat({block_m / warp_m, block_n / warp_n},
                                          true, false); // 16*Y x N (Y warp)
137
138
139
  return block_layout->Repeat({warp_m / 16, 1}, false, false);
}

140
141
142
Fragment makeGemmFragmentA(const int block_m, const int block_n,
                           const int block_k, const int warp_m,
                           const int warp_n, const int element_size) {
143
144
145
146
147
148
  // assume not transposed
  ICHECK(block_m % warp_m == 0);
  ICHECK(block_n % warp_n == 0);
  ICHECK(warp_m % 16 == 0);
  ICHECK(block_k % 16 == 0);
  // Only support 8-bit and 16-bit
149
150
  ICHECK(element_size == 8 || element_size == 16)
      << "element bitwidth=" << element_size;
151
152
  if (element_size == 8) {
    auto base_layout = makeGemmFragment8x16()->Repeat({2, 2}, false, false);
153
154
155
156
    auto warp_layout = base_layout->Repeat({block_m / warp_m, 1}, true)
                           ->Replicate(block_n / warp_n);
    auto block_layout =
        warp_layout->Repeat({warp_m / 16, block_k / 32}, false, false);
157
158
159
    return block_layout;
  } else if (element_size == 16) {
    auto base_layout = makeGemmFragment8x8()->Repeat({2, 2}, false, false);
160
161
162
163
    auto warp_layout = base_layout->Repeat({block_m / warp_m, 1}, true)
                           ->Replicate(block_n / warp_n);
    auto block_layout =
        warp_layout->Repeat({warp_m / 16, block_k / 16}, false, false);
164
165
166
167
168
169
170
    return block_layout;
  } else {
    ICHECK(0);
    return Fragment();
  }
}

171
172
Fragment makeGemmFragmentACDNA(const int block_m, const int block_n,
                               const int block_k, const int warp_m,
173
174
                               const int warp_n, const int element_size,
                               bool transposed) {
175
176
177
178
179
  // assume not transposed
  ICHECK(block_m % warp_m == 0);
  ICHECK(block_n % warp_n == 0);
  ICHECK(warp_m % 16 == 0);
  ICHECK(block_k % 16 == 0);
180
181
  ICHECK(element_size == 8 || element_size == 16)
      << "element bitwidth=" << element_size;
182
  if (transposed) {
183
184
185
    auto base_layout =
        makeGemmFragmentAB16x16CDNATransposed()->Repeat({1, 1}, false, false);
    auto warp_layout =
186
187
        base_layout->Repeat({block_k / 16, warp_m / 16}, false, true);
    auto block_layout = warp_layout->Repeat({1, block_m / warp_m}, true, true)
188
                            ->Replicate(block_n / warp_n);
189
190
    return block_layout;
  } else {
191
192
193
194
195
196
    auto base_layout =
        makeGemmFragmentAB16x16CDNA()->Repeat({1, 1}, false, false);
    auto warp_layout =
        base_layout->Repeat({warp_m / 16, block_k / 16}, false, false);
    auto block_layout = warp_layout->Repeat({block_m / warp_m, 1}, true, true)
                            ->Replicate(block_n / warp_n);
197
198
199
200
    return block_layout;
  }
}

201
202
203
Fragment makeGemmFragmentB(const int block_m, const int block_n,
                           const int block_k, const int warp_m,
                           const int warp_n) {
204
205
206
  // transposed
  ICHECK(warp_n % 8 == 0);
  ICHECK(block_k % 16 == 0);
207
208
209
210
211
212
  auto base_layout =
      makeGemmFragment8x8Transposed()->Repeat({2, 1}, false, false);
  auto warp_layout = base_layout->Replicate(block_m / warp_m)
                         ->Repeat({1, block_n / warp_n}, true);
  auto block_layout =
      warp_layout->Repeat({block_k / 16, warp_n / 8}, false, true);
213
214
215
216
217
218
219
220
221
222
223
  return block_layout;
}

Fragment makeGemmFragment32x32(int element_size) {
  IterVar i = make_itervar("i", 32);
  IterVar j = make_itervar("j", 32);
  IterVar rep = make_itervar("rep", 1);
  ICHECK(element_size == 16 || element_size == 32);
  if (element_size == 16) {
    PrimExpr thd = FloorMod(i, 4) + FloorDiv(FloorMod(i, 16), 8) * 4 +
                   FloorDiv(FloorMod(j, 16), 8) * 8 + FloorDiv(i, 16) * 16;
224
225
    PrimExpr idx = FloorMod(j, 4) + FloorDiv(j, 16) * 4 +
                   FloorDiv(FloorMod(i, 8), 4) * 8 +
226
227
228
229
                   FloorDiv(FloorMod(j, 8), 4) * 16;
    return Fragment({i, j}, {idx}, thd, rep);
  } else {
    PrimExpr thd = FloorMod(i, 2) + 2 * FloorDiv(FloorMod(j, 4), 2) +
230
231
232
233
234
                   FloorDiv(FloorMod(i, 16), 8) * 4 +
                   FloorDiv(FloorMod(j, 16), 8) * 8 + FloorDiv(i, 16) * 16;
    PrimExpr idx = FloorMod(j, 2) + 2 * FloorDiv(FloorMod(i, 4), 2) +
                   FloorDiv(j, 16) * 4 + FloorDiv(FloorMod(i, 8), 4) * 8 +
                   FloorDiv(FloorMod(j, 8), 4) * 16;
235
236
237
238
    return Fragment({i, j}, {idx}, thd, rep);
  }
}

239
240
241
Fragment makeGemmVoltaFragmentC(const int block_m, const int block_n,
                                const int warp_m, const int warp_n,
                                int element_size) {
242
243
244
245
246
  ICHECK(block_m % warp_m == 0);
  ICHECK(block_n % warp_n == 0);
  ICHECK(warp_m % 32 == 0);
  ICHECK(warp_n % 32 == 0);
  auto base_layout = makeGemmFragment32x32(element_size);
247
248
249
250
  auto warp_layout =
      base_layout->Repeat({warp_m / 32, warp_n / 32}, false, false);
  auto block_layout =
      warp_layout->Repeat({block_m / warp_m, block_n / warp_n}, true);
251
252
253
  return block_layout;
}

254
255
256
Fragment makeGemmVoltaFragmentA(const int block_m, const int block_n,
                                const int block_k, const int warp_m,
                                const int warp_n) {
257
258
259
260
261
262
263
264
265
  // assume not transposed
  ICHECK(block_m % warp_m == 0);
  ICHECK(block_n % warp_n == 0);
  ICHECK(warp_m % 32 == 0);
  ICHECK(block_k % 4 == 0);
  // this is a special case
  IterVar i = make_itervar("i", 32);
  IterVar j = make_itervar("j", 4);
  IterVar rep = make_itervar("rep", 2);
266
267
  PrimExpr thd = FloorDiv(FloorMod(i, 16), 8) * 4 + 16 * FloorDiv(i, 16) +
                 FloorMod(i, 4) + 8 * rep;
268
269
  PrimExpr idx = j + FloorDiv(FloorMod(i, 8), 4) * 4;
  Fragment base_layout = Fragment({i, j}, {idx}, thd, rep);
270
271
272
273
  auto warp_layout =
      base_layout->Repeat({warp_m / 32, block_k / 4}, false, false);
  auto block_layout = warp_layout->Replicate(block_n / warp_n)
                          ->Repeat({block_m / warp_m, 1}, true);
274
275
276
  return block_layout;
}

277
278
279
PrimExpr xor2x2(const PrimExpr &i, const PrimExpr &j) {
  return FloorMod(i + j, 2);
}
280

281
PrimExpr xor4x4(const PrimExpr &i, const PrimExpr &j) {
282
283
284
285
286
287
288
  PrimExpr i0 = FloorMod(i, 2);
  PrimExpr j0 = FloorMod(j, 2);
  PrimExpr i1 = FloorDiv(i, 2);
  PrimExpr j1 = FloorDiv(j, 2);
  return 2 * xor2x2(i1, j1) + xor2x2(i0, j0);
}

289
PrimExpr xor8x8(const PrimExpr &i, const PrimExpr j) {
290
291
292
293
294
295
296
297
298
299
300
301
  PrimExpr i0 = FloorMod(i, 2);
  PrimExpr j0 = FloorMod(j, 2);
  PrimExpr i1 = FloorDiv(i, 2);
  PrimExpr j1 = FloorDiv(j, 2);
  return 2 * xor4x4(i1, j1) + xor2x2(i0, j0);
}

Layout makeHalfBankSwizzleLayout(int stride, int continuous, int element_size) {
  // Swizzle 2 bit
  Var i = InputPlaceholder(0);
  Var j = InputPlaceholder(1);
  int vector_size = 128 / element_size;
302
303
304
  ICHECK(stride % 8 == 0) << "stride=" << stride;
  ICHECK(continuous % (vector_size * 4) == 0)
      << "continuous=" << continuous << ", vector_size=" << vector_size;
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
  PrimExpr ts = FloorDiv(i, 8);
  PrimExpr s = FloorMod(i, 8);
  PrimExpr tc = FloorDiv(FloorDiv(j, vector_size), 4);
  PrimExpr c = FloorMod(FloorDiv(j, vector_size), 4);
  PrimExpr vec = FloorMod(j, vector_size);
  PrimExpr c_swizzle = xor4x4(c, FloorDiv(s, 2));
  PrimExpr index = vec + (c_swizzle + s * 4) * vector_size;
  return Layout(Array<PrimExpr>{stride, continuous}, {tc, ts, index});
}

Layout makeFullBankSwizzleLayout(int stride, int continuous, int element_size) {
  // Swizzle 3 bit
  Var i = InputPlaceholder(0);
  Var j = InputPlaceholder(1);
  int vector_size = 128 / element_size;
  ICHECK(stride % 8 == 0);
  ICHECK(continuous % (vector_size * 8) == 0);
  PrimExpr ts = FloorDiv(i, 8);
  PrimExpr s = FloorMod(i, 8);
  PrimExpr tc = FloorDiv(FloorDiv(j, vector_size), 8);
  PrimExpr c = FloorMod(FloorDiv(j, vector_size), 8);
  PrimExpr vec = FloorMod(j, vector_size);
  PrimExpr c_swizzle = xor8x8(c, s);
  PrimExpr index = vec + (c_swizzle + s * 8) * vector_size;
  return Layout(Array<PrimExpr>{stride, continuous}, {tc, ts, index});
}

332
333
334
335
// Detail implementation please ref to
// bitblas::tl::mfma_layout::make_mfma_swizzle_layout
Layout makeMatrixCoreSwizzleLayout(int stride, int continuous, int element_size,
                                   int kPack = 1) {
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
  const int numBanks = 32;
  const int bankBitWidth = 32;
  const int SIMDWidth = 16;
  const int vecSize = 4 * kPack;
  const int innerDimLength = continuous;
  const int typeWidthInBit = element_size;

  const int elemsPerOneBanksRow = (numBanks * bankBitWidth) / typeWidthInBit;
  const int perPhase = std::max(1, elemsPerOneBanksRow / innerDimLength);
  const int maxPhase = std::min(SIMDWidth / perPhase, innerDimLength / vecSize);

  IterVar row = make_itervar("row", stride);
  IterVar col = make_itervar("col", continuous);
  PrimExpr phase = FloorMod(row / perPhase, maxPhase);
  PrimExpr colOffSwizzled = ((col / vecSize) ^ phase) * vecSize;
  PrimExpr colOffOrdered = FloorMod(col, vecSize);
  PrimExpr colOff = colOffSwizzled + colOffOrdered;

  return Layout(Array{row, col}, {row, colOff});
}

Layout makeGemmABLayoutF64_Kinner(int stride, int continuous) {
  // Swizzle<2, 0, 4>
  Var i = InputPlaceholder(0);
  Var j = InputPlaceholder(1);
  PrimExpr tc = FloorDiv(j, 16);
  PrimExpr ts = FloorDiv(i, 4);
  PrimExpr c = FloorMod(j, 16);
  PrimExpr s = FloorMod(i, 4);
  PrimExpr swizzled_c = FloorDiv(c, 4) * 4 + xor4x4(FloorMod(c, 4), s);
  PrimExpr index = swizzled_c + s * 16;
  return Layout(Array<PrimExpr>{stride, continuous}, {tc, ts, index});
}

Layout makeGemmABLayoutF64_Kouter(int stride, int continuous) {
  // Swizzle<2, 2, 2>
  Var i = InputPlaceholder(0);
  Var j = InputPlaceholder(1);
  PrimExpr tc = FloorDiv(j, 16);
  PrimExpr ts = FloorDiv(i, 4);
  PrimExpr c = FloorMod(j, 16);
  PrimExpr s = FloorMod(i, 4);
  PrimExpr swizzled_c = FloorMod(c, 4) + xor4x4(FloorDiv(c, 4), s) * 4;
  PrimExpr index = swizzled_c + s * 16;
  return Layout(Array<PrimExpr>{stride, continuous}, {tc, ts, index});
}

// The Default Layout for Tensor Access
Layout makeGemmLayoutLinear(int stride, int continuous) {
  IterVar i = make_itervar("i", stride);
  IterVar j = make_itervar("j", continuous);
  return Layout(Array{i, j}, {i * continuous + j});
}

Layout makeGemmABLayoutPadded(int stride, int continuous, int element_size) {
  IterVar i = make_itervar("i", stride);
  IterVar j = make_itervar("j", continuous);
  int padded = continuous;
  // Add 128 bits padding when the last dim is a multiple of 256 bits
395
396
  if ((element_size * continuous) % 256 == 0)
    padded += 128 / element_size;
397
398
399
400
401
402
403
404
405
406
  return Layout(Array{i, j}, {i * padded + j});
}

Layout MakeGemmVoltaABLayoutCrosswise(int stride, int continuous) {
  ICHECK(stride % 32 == 0 && continuous % 32 == 0);
  IterVar i = make_itervar("i", stride);
  IterVar j = make_itervar("j", continuous);
  PrimExpr vec_contiguous_idx = FloorDiv(j, 4);
  PrimExpr vec_strided_within_tile = FloorMod(vec_contiguous_idx, 8);

407
408
409
410
411
412
413
414
415
416
417
  PrimExpr bit2 =
      FloorMod(FloorDiv(FloorMod(i, 32), 16) + FloorDiv(FloorMod(i, 16), 8) +
                   FloorDiv(vec_strided_within_tile, 4),
               2);
  PrimExpr bit1 = xor2x2(FloorDiv(FloorMod(i, 8), 4),
                         FloorDiv(FloorMod(vec_strided_within_tile, 4), 2));
  PrimExpr permuted_vec_contiguous =
      FloorDiv(i, 16) * 16 + FloorMod(i, 4) * 4 + bit2 * 2 + bit1;

  PrimExpr offset = FloorMod(j, 4) + permuted_vec_contiguous * 4 +
                    vec_contiguous_idx * stride * 4;
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
  return Layout(Array{i, j}, {offset});
}

Layout MakeGemmVoltaALayoutCongruous(int stride, int continuous) {
  ICHECK(stride % 4 == 0 && continuous % 64 == 0);
  IterVar i = make_itervar("i", stride);
  IterVar j = make_itervar("j", continuous);
  PrimExpr vec_contiguous_idx = FloorDiv(j, 8);
  PrimExpr vec_strided_idx = i;
  PrimExpr tile_contiguous_idx = FloorDiv(vec_contiguous_idx, 8);
  PrimExpr tile_strided_idx = FloorDiv(vec_strided_idx, 4);
  PrimExpr tile_contiguous_residual = FloorMod(vec_contiguous_idx, 8);
  PrimExpr tile_strided_residual = FloorMod(vec_strided_idx, 4);

  PrimExpr permuted_strided_within_tile = FloorDiv(tile_contiguous_residual, 2);
  PrimExpr permuted_contiguous_within_tile =
      FloorMod(tile_contiguous_residual, 2) * 4 +
      xor4x4(tile_strided_residual, permuted_strided_within_tile);

437
438
  PrimExpr element_strided =
      permuted_strided_within_tile + tile_strided_idx * 4;
439
  PrimExpr element_contiguous =
440
441
      FloorMod(j, 8) +
      (permuted_contiguous_within_tile + tile_contiguous_idx * 8) * 8;
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
  PrimExpr offset = element_strided * continuous + element_contiguous;
  return Layout(Array{i, j}, {offset});
}

Layout MakeGemmVoltaBLayoutCongruous(int stride, int continuous) {
  ICHECK(stride % 4 == 0 && continuous % 64 == 0);
  IterVar i = make_itervar("i", stride);
  IterVar j = make_itervar("j", continuous);
  PrimExpr vec_contiguous_idx = FloorDiv(j, 8);
  PrimExpr vec_strided_idx = i;
  PrimExpr tile_contiguous_idx = FloorDiv(vec_contiguous_idx, 8);
  PrimExpr tile_strided_idx = FloorDiv(vec_strided_idx, 4);
  PrimExpr tile_contiguous_residual = FloorMod(vec_contiguous_idx, 8);
  PrimExpr tile_strided_residual = FloorMod(vec_strided_idx, 4);

  PrimExpr permuted_strided_within_tile = FloorMod(tile_contiguous_residual, 4);
  PrimExpr permuted_contiguous_within_tile =
      FloorDiv(tile_contiguous_residual, 4) * 4 +
      xor4x4(tile_strided_residual, permuted_strided_within_tile);

462
463
  PrimExpr element_strided =
      permuted_strided_within_tile + tile_strided_idx * 4;
464
  PrimExpr element_contiguous =
465
466
      FloorMod(j, 8) +
      (permuted_contiguous_within_tile + tile_contiguous_idx * 8) * 8;
467
468
469
470
  PrimExpr offset = element_strided * continuous + element_contiguous;
  return Layout(Array{i, j}, {offset});
}

471
472
473
474
475
476
477
478
Layout makeGemmVoltaABLayout(int stride, int continuous, bool is_a,
                             int kfactor) {
  if (kfactor == 2)
    return MakeGemmVoltaABLayoutCrosswise(stride, continuous);
  if (is_a && continuous % 64 == 0)
    return MakeGemmVoltaALayoutCongruous(stride, continuous);
  if (!is_a && continuous % 64 == 0)
    return MakeGemmVoltaBLayoutCongruous(stride, continuous);
479
480
481
  return makeGemmABLayoutPadded(stride, continuous, 16);
}

482
483
Layout makeGemmABLayout(int mat_stride, int mat_continuous, int continuity,
                        int element_size, int kfactor) {
484
  if (element_size == 64) {
485
486
487
488
489
    if (kfactor == 1 && continuity % 16 == 0) // float64 KxN
      return makeGemmABLayoutF64_Kouter(mat_stride, mat_continuous);
    if (kfactor == 2 && continuity % 16 == 0) // float64 NxK
      return makeGemmABLayoutF64_Kinner(mat_stride, mat_continuous);
    return makeGemmABLayoutPadded(mat_stride, mat_continuous, element_size);
490
491
  }
  int vector_size = 128 / element_size;
492
  if (kfactor == 1 && element_size == 8) // int8 KxN
493
494
495
496
497
    return makeGemmABLayoutPadded(mat_stride, mat_continuous, element_size);
  else if (continuity % (vector_size * 8) == 0)
    return makeFullBankSwizzleLayout(mat_stride, mat_continuous, element_size);
  else if (continuity % (vector_size * 4) == 0)
    return makeHalfBankSwizzleLayout(mat_stride, mat_continuous, element_size);
498
  else {
499
    return makeGemmABLayoutPadded(mat_stride, mat_continuous, element_size);
500
501
502
  }
}

503
504
Layout makeGemmABLayoutCDNA(int stride, int continuous, int element_size,
                            int kPack) {
505
506
507
508
509
510
511
  int vector_size = 128 / element_size;
  if (continuous % (vector_size * 4) == 0)
    return makeMatrixCoreSwizzleLayout(stride, continuous, element_size, kPack);
  else {
    return makeGemmABLayoutPadded(stride, continuous, element_size);
  }
}
512
513
} // namespace tl
} // namespace tvm