builtin.h 16.9 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

namespace attr {
25
static constexpr const char *kSafeValueMap = "safe_value_map";
26
27
static constexpr const char *kWarpSpecializationScope =
    "kWarpSpecializationScope";
28
29
static constexpr const char *kCustomWarpSpecialization =
    "kCustomWarpSpecialization";
30
static constexpr const char *kLocalVarInit = "tl.local_var_init";
31
32
33
34
// A PrimFunc-level attribute carrying a list of handle Vars
// that must NOT be marked with the restrict qualifier in codegen.
// Type: Array<tir::Var>
static constexpr const char *kNonRestrictParams = "tl.non_restrict_params";
35
36
} // namespace attr

37
38
static constexpr const char *kDebugMergeSharedMemoryAllocations =
    "tl.debug_merge_shared_memory_allocations";
39
static constexpr const char *kDisableTMALower = "tl.disable_tma_lower";
40
41
static constexpr const char *kDisableSafeMemoryLegalize =
    "tl.disable_safe_memory_legalize";
42
43
static constexpr const char *kDisableWarpSpecialized =
    "tl.disable_warp_specialized";
44
static constexpr const char *kConfigIndexBitwidth = "tl.config_index_bitwidth";
45
46
static constexpr const char *kEnableAggressiveSharedMemoryMerge =
    "tl.enable_aggressive_shared_memory_merge";
47
static constexpr const char *kDisableFastMath = "tl.disable_fast_math";
48
static constexpr const char *kEnableFastMath = "tl.enable_fast_math";
49
50
static constexpr const char *kPtxasRegisterUsageLevel =
    "tl.ptxas_register_usage_level";
51
52
static constexpr const char *kEnablePTXASVerboseOutput =
    "tl.enable_ptxas_verbose_output";
53
static constexpr const char *kDisableVectorize256 = "tl.disable_vectorize_256";
54
static constexpr const char *kDisableWGMMA = "tl.disable_wgmma";
55
static constexpr const char *kDisableShuffleElect = "tl.disable_shuffle_elect";
56
57
static constexpr const char *kStorageRewriteDetectInplace =
    "tl.storage_rewrite_detect_inplace";
58
59
60
61
static constexpr const char *kLayoutVisualizationEnable =
    "tl.layout_visualization_enable";
static constexpr const char *kLayoutVisualizationFormats =
    "tl.layout_visualization_formats";
62
static constexpr const char *kDeviceCompileFlags = "tl.device_compile_flags";
63
64
65
66
67
68
69
70
71
/*!
 * \brief Whether to disable dynamic tail split
 *
 * kDisableDynamicTailSplit = "tl.disable_dynamic_tail_split"
 *
 */
static constexpr const char *kDisableDynamicTailSplit =
    "tl.disable_dynamic_tail_split";

72
73
74
75
76
77
78
79
80
81
82
83
84
85
/*!
 * \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";

86
87
88
89
90
91
92
93
/*!
 * \brief Force inline Let bindings during simplification.
 *
 * kForceLetInline = "tl.force_let_inline"
 *
 */
static constexpr const char *kForceLetInline = "tl.force_let_inline";

94
95
96
97
98
99
100
101
102
103
104
105
/*!
 * \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";

106
107
108
109
110
111
112
113
/*!
 * \brief Get the type of the CUDA tensor map
 *
 * DataType cuTensorMapType()
 *
 */
DataType cuTensorMapType();

114
// fast math related op
115
// __exp(x) - fast exponential
116
TVM_DLL const Op &__exp();
117
// __exp10(x) - fast base-10 exponential
118
TVM_DLL const Op &__exp10();
119
// __log(x) - fast natural logarithm
120
TVM_DLL const Op &__log();
121
// __log2(x) - fast base-2 logarithm
122
TVM_DLL const Op &__log2();
123
// __log10(x) - fast base-10 logarithm
124
TVM_DLL const Op &__log10();
125
// __tan(x) - fast tangent
126
TVM_DLL const Op &__tan();
127
// __cos(x) - fast cosine
128
TVM_DLL const Op &__cos();
129
// __sin(x) - fast sine
130
131
TVM_DLL const Op &__sin();

132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
// 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();

150
151
152
/*!
 * \brief tvm intrinsics for TMADescriptor creation for tiled load
 *
153
 * CuTensorMap* create_tma_descriptor(data_type, rank, global_addr,
154
155
 * global_shape..., global_stride..., smem_box..., smem_stride..., interleave,
 * swizzle, l2_promotion, oob_fill)
156
157
 *
 */
158
TVM_DLL const Op &create_tma_descriptor();
159
160
161
162

/*!
 * \brief tvm intrinsics for TMADescriptor creation for image to column load
 *
163
 * CuTensorMap* create_tma_im2col_descriptor(data_type, rank, global_addr,
164
165
166
 * global_shape..., global_stride..., elem_stride..., lower_corner...,
 * upper_corner..., smme_box_pixel, smem_box_channel, interleave, swizzle,
 * l2_promotion, oob_fill)
167
168
 *
 */
169
TVM_DLL const Op &create_tma_im2col_descriptor();
170
171
172
173

/*!
 * \brief Create a list of mbarrier with num_threads
 *
174
 * create_list_of_mbarrier(num_threads0, num_threads1, ...)
175
176
 *
 */
177
TVM_DLL const Op &create_list_of_mbarrier();
178
179
180
181
182
183
184

/*!
 * \brief Get the mbarrier with barrier_id
 *
 * int64_t* GetMBarrier(barrier_id)
 *
 */
185
TVM_DLL const Op &get_mbarrier();
186
187

/*!
188
189
 * \brief tvm intrinsics for loading data from global tensor descriptor to
 * shared memory
190
 *
191
 * tma_load(descriptor, mbarrier, smem_data, coord_0, coord_1, ...)
192
193
 *
 */
194
TVM_DLL const Op &tma_load();
195
196

/*!
197
198
 * \brief tvm intrinsics for loading image from global tensor to columns in
 * shared memory
199
 *
200
 * tma_load(descriptor, mbarrier, smem_data, coord_0, coord_1, ...,
201
 * image_offset, ...)
202
203
 *
 */
204
TVM_DLL const Op &tma_load_im2col();
205
206

/*!
207
208
 * \brief tvm intrinsics for storing data from shared memory to global tensor
 * descriptor
209
 *
210
 * tma_store(descriptor, smem_data, coord_0, coord_1, ...)
211
212
 *
 */
213
TVM_DLL const Op &tma_store();
214

215
216
217
218
219
220
221
222
/*!
 * \brief tvm intrinsics for barrier initialization fence
 *
 * ptx_fence_barrier_init()
 *
 */
const Op &ptx_fence_barrier_init();

223
224
225
/*!
 * \brief tvm intrinsics for mbarrier wait with parity bit
 *
226
 * mbarrier_wait_parity(mbarrier, parity)
227
228
 *
 */
229
TVM_DLL const Op &mbarrier_wait_parity();
230
231
232
233

/*!
 * \brief tvm intrinsics for mbarrier expect tx
 *
234
 * mbarrier_expect_tx(mbarrier, transaction_bytes)
235
236
 *
 */
237
TVM_DLL const Op &mbarrier_expect_tx();
238

239
240
241
242
243
244
245
246
247
248
249
250
251
252
/*!
 * \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.
 *
253
254
255
256
257
 *  void ptx_wgmma_rs(StringImm accum_dtype, StringImm wgmma_prefix,
 * 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);
258
259
260
 */
TVM_DLL const Op &ptx_wgmma_rs();

261
262
263
264
265
266
267
268
269
270
/*!
 * \brief tvm intrinsic for tcgen05 mma shared-shared instructions.
 */
TVM_DLL const Op &ptx_tcgen05_mma_ss();

/*!
 * \brief tvm intrinsic for tcgen05 mma tensor-shared instructions.
 */
TVM_DLL const Op &ptx_tcgen05_mma_ts();

271
272
273
274
275
276
/*!
 * \brief tvm intrinsics for initializing tensor memory
 *
 * ptx_init_tensor_memory(tmem_buffer, num_cols)
 *
 */
277
TVM_DLL const Op &ptx_init_tensor_memory();
278
279
280
281
282
283
284

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

287
288
289
290
291
292
293
294
295
296
297
/*!
 * \brief tvm intrinsic for ptx tensor core mma instructions on SM70.
 *
 *  void ptx_mma_sm70(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, bool saturate);
 */
TVM_DLL const Op &ptx_mma_sm70();

298
299
300
/*!
 * \brief tvm intrinsics for ldmatrix
 *
301
 * ptx_ldmatrix(transposed, num, shared_addr, local_addr)
302
303
 *
 */
304
TVM_DLL const Op &ptx_ldmatrix();
305
306
307
308

/*!
 * \brief tvm intrinsics for stmatrix
 *
309
 * ptx_ldmatrix(transposed, num, shared_addr, int32_values...)
310
311
 *
 */
312
TVM_DLL const Op &ptx_stmatrix();
313

314
315
316
317
318
319
320
321
/*!
 * \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();

322
323
324
/*!
 * \brief Pack two b16 value into a b32 value
 *
325
 * int32 pack_b16(b16_value, b16_value)
326
327
 *
 */
328
TVM_DLL const Op &pack_b16();
329
330
331
332
333
334
335

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

338
339
340
/*!
 * \brief Indicate arrival of warp issuing TMA_STORE
 *
341
 * tma_store_arrive()
342
343
 *
 */
344
TVM_DLL const Op &tma_store_arrive();
345
346
347
348

/*!
 * \brief Wait for TMA_STORE to finish
 *
349
 * tma_store_wait()
350
351
 *
 */
352
TVM_DLL const Op &tma_store_wait();
353

354
355
356
357
358
359
/*!
 * \brief Set reg hint for warp-specialized branched
 *
 * SetMaxNRegInc(num_reg, is_inc)
 *
 */
360
TVM_DLL const Op &set_max_nreg();
361

362
363
364
/*!
 * \brief No set reg hint for warp-specialized branched
 *
365
 * no_set_max_nreg()
366
367
 *
 */
368
TVM_DLL const Op &no_set_max_nreg();
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
/*!
 * \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();

394
395
396
397
398
399
400
401
/*!
 * \brief Fence accumulator operand registers for upcoming WGMMA operations
 *
 * warpgroup_fence_operand(dtype, ptr, offset, num_regs)
 *
 */
TVM_DLL const Op &warpgroup_fence_operand();

402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
/*!
 * \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();

434
435
436
/*!
 * \brief Wait the previous wgmma to finish
 *
437
 * wait_wgmma(num_mma)
438
439
 *
 */
440
TVM_DLL const Op &wait_wgmma();
441

442
443
444
445
446
447
/*!
 * \brief Synchronize all threads in a grid
 *
 * sync_grid()
 *
 */
448
TVM_DLL const Op &sync_grid();
449
450
451
452
453
454
455

/*!
 * \brief tvm intrinsic for loop continue
 *
 * loop_break()
 *
 */
456
TVM_DLL const Op &loop_break();
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
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
/*!
 * \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();

506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
/*!
 * \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();

521
522
523
524
525
526
527
/*!
 * \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();

528
529
530
531
532
533
534
/*!
 * \brief tilelang intrinsic for initializing a descriptor buffer for
 * wgmma/utcmma.
 *
 *  This op is used to represent a descriptor initialization operation in
 * tilelang.
 */
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
TVM_DLL const Op &initialize_wgmma_descriptor();

/*!
 * \brief tilelang intrinsic for initializing a descriptor buffer for
 * tcgen05 mma.
 */
TVM_DLL const Op &initialize_tcgen05_descriptor();

/*!
 * \brief tilelang intrinsic for committing UMMA (TCGEN05) barrier arrive.
 *
 *  This op wraps the device-side arrive used to signal completion of MMA work
 *  to a shared-memory mbarrier. It mirrors CUTLASS's umma_arrive.
 */
TVM_DLL const Op &tcgen05_mma_arrive();
550
551
552
553
554
555
556
557

/*!
 * \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.
 */
558

559
TVM_DLL const Op &increase_descriptor_offset();
560

561
562
563
564
565
566
567
/*!
 * \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();
568

569
570
571
572
573
574
575
576
577
578
579
580
581
582
/*!
 * \brief tilelang intrinsic for assert on device.
 *
 *  This op is used to represent an assert on device
 */
TVM_DLL const Op &device_assert();

/*!
 * \brief tilelang intrinsic for assert on device with additional message.
 *
 *  This op is used to represent an assert on device with additional message.
 */
TVM_DLL const Op &device_assert_with_msg();

Tong WU's avatar
Tong WU committed
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
/*!
 * \brief tilelang intrinsic for warp reduction sum.
 */
TVM_DLL const Op &warp_reduce_sum();

/*!
 * \brief tilelang intrinsic for warp reduction max.
 */
TVM_DLL const Op &warp_reduce_max();

/*!
 * \brief tilelang intrinsic for warp reduction min.
 */
TVM_DLL const Op &warp_reduce_min();

/*!
 * \brief tilelang intrinsic for warp reduction bitand.
 */
TVM_DLL const Op &warp_reduce_bitand();

/*!
 * \brief tilelang intrinsic for warp reduction bitor.
 */
TVM_DLL const Op &warp_reduce_bitor();

608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
/*!
 * \brief tilelang intrinsic for CUDA read-only cache load (__ldg).
 *
 *  This op allows users to explicitly request a non-coherent cached load
 *  from global memory on CUDA by emitting `__ldg(&ptr[idx])` for 32-bit
 *  element types on supported architectures. It provides a direct way to
 *  leverage the read-only data cache for performance-sensitive loads when
 *  the compiler cannot infer `const __restrict__` automatically.
 *
 *  Usage from TVMScript:
 *    y[i] = T.__ldg(x[i])
 *
 *  The op takes one argument preferred as a BufferLoad identifying the
 *  source element; alternatively, backends may support passing a Buffer and
 *  index expression.
 */
TVM_DLL const Op &__ldg();

626
627
} // namespace tl
} // namespace tvm
628

629
#endif //  TVM_TL_OP_BUILTIN_H_