builtin.h 13.4 KB
Newer Older
1
2
3
4
5
6
7
8
9
/*!
 * \file tl/op/builtin.h
 * \brief Builtin intrinsics.
 *
 */

#ifndef TVM_TL_OP_BUILTIN_H_
#define TVM_TL_OP_BUILTIN_H_

10
#include "operator.h"
11
#include <tvm/ir/transform.h>
12
13

namespace tvm {
14
15
16
17
18
19
20
21
/*!
 * \brief Create the TVM intrinsic that initializes a PTX fence barrier.
 *
 * Initializes a PTX fence-style barrier used to coordinate asynchronous memory
 * operations (for example, TMA/TMA_STORE). Returns the Op representing this
 * intrinsic for use in TIR lowering and code generation.
 *
 */
22
namespace tl {
23
24
25

namespace attr {
static constexpr const char *kPaddingMap = "padding_map";
26
27
static constexpr const char *kWarpSpecializationScope =
    "kWarpSpecializationScope";
28
29
static constexpr const char *kCustomWarpSpecialization =
    "kCustomWarpSpecialization";
30
31
} // namespace attr

32
33
static constexpr const char *kDebugMergeSharedMemoryAllocations =
    "tl.debug_merge_shared_memory_allocations";
34
static constexpr const char *kDisableTMALower = "tl.disable_tma_lower";
35
36
static constexpr const char *kDisableSafeMemoryLegalize =
    "tl.disable_safe_memory_legalize";
37
38
static constexpr const char *kDisableWarpSpecialized =
    "tl.disable_warp_specialized";
39
static constexpr const char *kConfigIndexBitwidth = "tl.config_index_bitwidth";
40
41
static constexpr const char *kEnableAggressiveSharedMemoryMerge =
    "tl.enable_aggressive_shared_memory_merge";
42
static constexpr const char *kDisableFastMath = "tl.disable_fast_math";
43
static constexpr const char *kEnableFastMath = "tl.enable_fast_math";
44
45
static constexpr const char *kPtxasRegisterUsageLevel =
    "tl.ptxas_register_usage_level";
46
47
static constexpr const char *kEnablePTXASVerboseOutput =
    "tl.enable_ptxas_verbose_output";
48
static constexpr const char *kDisableVectorize256 = "tl.disable_vectorize_256";
49
static constexpr const char *kDisableWGMMA = "tl.disable_wgmma";
50
static constexpr const char *kDisableShuffleElect = "tl.disable_shuffle_elect";
51
52
53
54
55
56
57
58
59
/*!
 * \brief Whether to disable dynamic tail split
 *
 * kDisableDynamicTailSplit = "tl.disable_dynamic_tail_split"
 *
 */
static constexpr const char *kDisableDynamicTailSplit =
    "tl.disable_dynamic_tail_split";

60
61
62
63
64
65
66
67
68
69
70
71
72
73
/*!
 * \brief Whether to disable thread storage synchronization
 *
 * When enabled, disables the automatic insertion of thread synchronization
 * barriers (e.g., __syncthreads()) for shared memory access coordination.
 * This can be useful for performance optimization in cases where manual
 * synchronization is preferred or when synchronization is not needed.
 *
 * kDisableThreadStorageSync = "tl.disable_thread_storage_sync"
 *
 */
static constexpr const char *kDisableThreadStorageSync =
    "tl.disable_thread_storage_sync";

74
75
76
77
78
79
80
81
/*!
 * \brief Force inline Let bindings during simplification.
 *
 * kForceLetInline = "tl.force_let_inline"
 *
 */
static constexpr const char *kForceLetInline = "tl.force_let_inline";

82
83
84
85
86
87
88
89
90
91
92
93
/*!
 * \brief The size of the vectorized dimension in buffer, designed by user
 *
 * For example, if the vectorized dimension is 128 bits and the dtype of buffer
 * A[m, k] is float16, the size of the vectorized dimension (i.e. k) in buffer A
 * should be divisible by 8 (8 = 128 / 16).
 *
 * kDynamicAlignment = "tl.dynamic_alignment"
 *
 */
static constexpr const char *kDynamicAlignment = "tl.dynamic_alignment";

94
95
96
97
98
99
100
101
/*!
 * \brief Get the type of the CUDA tensor map
 *
 * DataType cuTensorMapType()
 *
 */
DataType cuTensorMapType();

102
// fast math related op
103
// __exp(x) - fast exponential
104
TVM_DLL const Op &__exp();
105
// __exp10(x) - fast base-10 exponential
106
TVM_DLL const Op &__exp10();
107
// __log(x) - fast natural logarithm
108
TVM_DLL const Op &__log();
109
// __log2(x) - fast base-2 logarithm
110
TVM_DLL const Op &__log2();
111
// __log10(x) - fast base-10 logarithm
112
TVM_DLL const Op &__log10();
113
// __tan(x) - fast tangent
114
TVM_DLL const Op &__tan();
115
// __cos(x) - fast cosine
116
TVM_DLL const Op &__cos();
117
// __sin(x) - fast sine
118
119
TVM_DLL const Op &__sin();

120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
// high precision with IEEE-compliant.
// ieee_add(x, y, rounding_mode) - IEEE-compliant addition
TVM_DLL const Op &ieee_add();
// ieee_sub(x, y, rounding_mode) - IEEE-compliant subtraction
TVM_DLL const Op &ieee_sub();
// ieee_mul(x, y, rounding_mode) - IEEE-compliant multiplication
TVM_DLL const Op &ieee_mul();
// ieee_fmaf(x, y, z, rounding_mode) - IEEE-compliant fused multiply-add
TVM_DLL const Op &ieee_fmaf();
// ieee_frcp(x, rounding_mode) - IEEE-compliant reciprocal
TVM_DLL const Op &ieee_frcp();
// ieee_fsqrt(x, rounding_mode) - IEEE-compliant square root
TVM_DLL const Op &ieee_fsqrt();
// ieee_frsqrt(x) - IEEE-compliant reciprocal square root (rn only)
TVM_DLL const Op &ieee_frsqrt();
// ieee_fdiv(x, y, rounding_mode) - IEEE-compliant division
TVM_DLL const Op &ieee_fdiv();

138
139
140
/*!
 * \brief tvm intrinsics for TMADescriptor creation for tiled load
 *
141
 * CuTensorMap* create_tma_descriptor(data_type, rank, global_addr,
142
143
 * global_shape..., global_stride..., smem_box..., smem_stride..., interleave,
 * swizzle, l2_promotion, oob_fill)
144
145
 *
 */
146
TVM_DLL const Op &create_tma_descriptor();
147
148
149
150

/*!
 * \brief tvm intrinsics for TMADescriptor creation for image to column load
 *
151
 * CuTensorMap* create_tma_im2col_descriptor(data_type, rank, global_addr,
152
153
154
 * global_shape..., global_stride..., elem_stride..., lower_corner...,
 * upper_corner..., smme_box_pixel, smem_box_channel, interleave, swizzle,
 * l2_promotion, oob_fill)
155
156
 *
 */
157
TVM_DLL const Op &create_tma_im2col_descriptor();
158
159
160
161

/*!
 * \brief Create a list of mbarrier with num_threads
 *
162
 * create_list_of_mbarrier(num_threads0, num_threads1, ...)
163
164
 *
 */
165
TVM_DLL const Op &create_list_of_mbarrier();
166
167
168
169
170
171
172

/*!
 * \brief Get the mbarrier with barrier_id
 *
 * int64_t* GetMBarrier(barrier_id)
 *
 */
173
TVM_DLL const Op &get_mbarrier();
174
175

/*!
176
177
 * \brief tvm intrinsics for loading data from global tensor descriptor to
 * shared memory
178
 *
179
 * tma_load(descriptor, mbarrier, smem_data, coord_0, coord_1, ...)
180
181
 *
 */
182
TVM_DLL const Op &tma_load();
183
184

/*!
185
186
 * \brief tvm intrinsics for loading image from global tensor to columns in
 * shared memory
187
 *
188
 * tma_load(descriptor, mbarrier, smem_data, coord_0, coord_1, ...,
189
 * image_offset, ...)
190
191
 *
 */
192
TVM_DLL const Op &tma_load_im2col();
193
194

/*!
195
196
 * \brief tvm intrinsics for storing data from shared memory to global tensor
 * descriptor
197
 *
198
 * tma_store(descriptor, smem_data, coord_0, coord_1, ...)
199
200
 *
 */
201
TVM_DLL const Op &tma_store();
202

203
204
205
206
207
208
209
210
/*!
 * \brief tvm intrinsics for barrier initialization fence
 *
 * ptx_fence_barrier_init()
 *
 */
const Op &ptx_fence_barrier_init();

211
212
213
/*!
 * \brief tvm intrinsics for mbarrier wait with parity bit
 *
214
 * mbarrier_wait_parity(mbarrier, parity)
215
216
 *
 */
217
TVM_DLL const Op &mbarrier_wait_parity();
218
219
220
221

/*!
 * \brief tvm intrinsics for mbarrier expect tx
 *
222
 * mbarrier_expect_tx(mbarrier, transaction_bytes)
223
224
 *
 */
225
TVM_DLL const Op &mbarrier_expect_tx();
226

227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
/*!
 * \brief tvm intrinsic for ptx tensor core wgmma instructions.
 *
 *  void ptx_wgmma_ss(StringImm accum_dtype, StringImm wgmma_prefix, bool
 * a_is_k_major, bool b_is_k_major, StringImm a_dtype_abbrv, StringImm
 * b_dtype_abbrv, StringImm accum_dtype_abbrv, Var A_descriptor, PrimExpr
 * A_offset, Var B_descriptor, Var B_offset, Var C_data, Var C_offset, bool
 * scale_out, bool scale_in_a, bool scale_in_b);
 */
TVM_DLL const Op &ptx_wgmma_ss();

/*!
 * \brief tvm intrinsics for ptx tensor core wgmma instructions.
 *
 *  void ptx_wgmma_rs(StringImm accum_dtype, StringImm wgmma_prefix, bool
 * a_is_k_major, bool b_is_k_major, StringImm a_dtype_abbrv, StringImm
 * b_dtype_abbrv, StringImm accum_dtype_abbrv, Var A_descriptor, PrimExpr
 * A_offset, Var B_descriptor, Var B_offset, Var C_data, Var C_offset, bool
 * scale_out, bool scale_in_a, bool scale_in_b);
 */
TVM_DLL const Op &ptx_wgmma_rs();

249
250
251
252
253
254
/*!
 * \brief tvm intrinsics for initializing tensor memory
 *
 * ptx_init_tensor_memory(tmem_buffer, num_cols)
 *
 */
255
TVM_DLL const Op &ptx_init_tensor_memory();
256
257
258
259
260
261
262

/*!
 * \brief tvm intrinsics for deallocating tensor memory
 *
 * tmem_deallocate(tmem_buffer)
 *
 */
263
TVM_DLL const Op &ptx_deallocate_tensor_memory();
264

265
266
267
/*!
 * \brief tvm intrinsics for ldmatrix
 *
268
 * ptx_ldmatrix(transposed, num, shared_addr, local_addr)
269
270
 *
 */
271
TVM_DLL const Op &ptx_ldmatrix();
272
273
274
275

/*!
 * \brief tvm intrinsics for stmatrix
 *
276
 * ptx_ldmatrix(transposed, num, shared_addr, int32_values...)
277
278
 *
 */
279
TVM_DLL const Op &ptx_stmatrix();
280

281
282
283
284
285
286
287
288
/*!
 * \brief tvm intrinsic for ptx async copy barrier using
 * cp.async.mbarrier.arrive.noinc
 *
 *  This op is used to represent a ptx async copy barrier operation in tilelang.
 */
TVM_DLL const Op &ptx_cp_async_barrier_noinc();

289
290
291
/*!
 * \brief Pack two b16 value into a b32 value
 *
292
 * int32 pack_b16(b16_value, b16_value)
293
294
 *
 */
295
TVM_DLL const Op &pack_b16();
296
297
298
299
300
301
302

/*!
 * \brief Issue a shared memory fence for async operations
 *
 * FenceProxyAsync()
 *
 */
303
TVM_DLL const Op &fence_proxy_async();
304

305
306
307
/*!
 * \brief Indicate arrival of warp issuing TMA_STORE
 *
308
 * tma_store_arrive()
309
310
 *
 */
311
TVM_DLL const Op &tma_store_arrive();
312
313
314
315

/*!
 * \brief Wait for TMA_STORE to finish
 *
316
 * tma_store_wait()
317
318
 *
 */
319
TVM_DLL const Op &tma_store_wait();
320

321
322
323
324
325
326
/*!
 * \brief Set reg hint for warp-specialized branched
 *
 * SetMaxNRegInc(num_reg, is_inc)
 *
 */
327
TVM_DLL const Op &set_max_nreg();
328

329
330
331
/*!
 * \brief No set reg hint for warp-specialized branched
 *
332
 * no_set_max_nreg()
333
334
 *
 */
335
TVM_DLL const Op &no_set_max_nreg();
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
/*!
 * \brief Arrive at a warpgroup fence for WGMMA sequences
 *
 * warpgroup_arrive()
 *
 */
TVM_DLL const Op &warpgroup_arrive();

/*!
 * \brief Commit the current warpgroup batch for WGMMA sequences
 *
 * warpgroup_commit_batch()
 *
 */
TVM_DLL const Op &warpgroup_commit_batch();

/*!
 * \brief Wait for the warpgroup batch identified by num_mma
 *
 * warpgroup_wait(num_mma)
 *
 */
TVM_DLL const Op &warpgroup_wait();

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
/*!
 * \brief Return the canonical lane index for the calling thread.
 *
 * get_lane_idx([warp_size])
 *
 */
TVM_DLL const Op &get_lane_idx();

/*!
 * \brief Return the canonical warp index, assuming converged threads.
 *
 * get_warp_idx_sync([warp_size])
 *
 */
TVM_DLL const Op &get_warp_idx_sync();

/*!
 * \brief Return the canonical warp index without synchronizing the warp.
 *
 * get_warp_idx([warp_size])
 *
 */
TVM_DLL const Op &get_warp_idx();

/*!
 * \brief Return the canonical warp group index for converged threads.
 *
 * get_warp_group_idx([warp_size, warps_per_group])
 *
 */
TVM_DLL const Op &get_warp_group_idx();

393
394
395
/*!
 * \brief Wait the previous wgmma to finish
 *
396
 * wait_wgmma(num_mma)
397
398
 *
 */
399
TVM_DLL const Op &wait_wgmma();
400

401
402
403
404
405
406
/*!
 * \brief Synchronize all threads in a grid
 *
 * sync_grid()
 *
 */
407
TVM_DLL const Op &sync_grid();
408
409
410
411
412
413
414

/*!
 * \brief tvm intrinsic for loop continue
 *
 * loop_break()
 *
 */
415
TVM_DLL const Op &loop_break();
416

417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
/*!
 * \brief tvm intrinsic for amd matrix core mfma instructions.
 *
 *  void tvm_mfma(StringImm shape, StringImm A_layout, StringImm B_layout,
 *               StringImm A_dtype, StringImm B_dtype, StringImm C_dtype,
 *               Var multiplicand_a, Expr a_index,
 *               Var multiplicand_b, Expr b_index,
 *               Var accumulator, Expr c_index);
 */
TVM_DLL const Op &tvm_mfma();

/*!
 * \brief tvm intrinsic for storing the result of AMD MFMA into a destination
 * pointer.
 *
 *        There is no real instruction that does that, but we want to hide
 * details of complex index manipulation behind this intrinsic to simplify TIR
 * lowering passes (e.g. LowerWarpMemory) like cuda ptx backend does.
 *
 * void tvm_mfma_store(IntImm m, IntImm n, Var dst_ptr, Var src_ptr, Expr
 * src_offset, Var dst_stride);
 */
TVM_DLL const Op &tvm_mfma_store();

/*!
 * \brief tvm intrinsic for amd rdna matrix core instructions.
 *
 *  void tvm_rdna_wmma(StringImm shape, StringImm A_layout, StringImm B_layout,
 *               StringImm A_dtype, StringImm B_dtype, StringImm C_dtype,
 *               Var multiplicand_a, Expr a_index,
 *               Var multiplicand_b, Expr b_index,
 *               Var accumulator, Expr c_index);
 */
TVM_DLL const Op &tvm_rdna_wmma();

/*!
 * \brief tvm intrinsic for storing the result of AMD RDNA WMMA into a
 * destination pointer.
 *
 *        There is no real instruction that does that, but we want to hide
 * details of complex index manipulation behind this intrinsic to simplify TIR
 * lowering passes (e.g. LowerWarpMemory) like cuda ptx backend does.
 *
 * void tvm_rdna_wmma_store(IntImm m, IntImm n, Var dst_ptr, Var src_ptr, Expr
 * src_offset, Var dst_stride);
 */
TVM_DLL const Op &tvm_rdna_wmma_store();

465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
/*!
 * \brief tilelang intrinsic for general matrix multiplication (GEMM).
 *
 *  This op is used to represent a generic GEMM operation in tilelang.
 */
TVM_DLL const Op &tl_gemm();

/*!
 * \brief tilelang intrinsic for sparse matrix multiplication (GEMM with
 * sparsity).
 *
 *  This op is used to represent a sparse GEMM operation in tilelang.
 */
TVM_DLL const Op &tl_gemm_sp();

480
481
482
483
484
485
486
/*!
 * \brief tilelang intrinsic for shuffle elect.
 *
 *  This op is used to represent a shuffle elect operation in tilelang.
 */
TVM_DLL const Op &tl_shuffle_elect();

487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
/*!
 * \brief tilelang intrinsic for initializing a descriptor buffer for
 * wgmma/utcmma.
 *
 *  This op is used to represent a descriptor initialization operation in
 * tilelang.
 */
TVM_DLL const Op &initialize_descriptor();

/*!
 * \brief tilelang intrinsic for setting the start address of a descriptor
 * buffer for wgmma/utcmma.
 *
 *  This op is used to represent a descriptor start address setting operation in
 * tilelang.
 */
TVM_DLL const Op &increase_descriptor_offset();
504
505
506
507
508
509
510
/*!
 * \brief tilelang intrinsic for element-wise atomic addition.
 *
 *  This op is used to represent an element-wise atomic add operation in
 * tilelang.
 */
TVM_DLL const Op &atomicadd_elem_op();
511

512
513
} // namespace tl
} // namespace tvm
514

515
#endif //  TVM_TL_OP_BUILTIN_H_