builtin.h 8.35 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
44
static constexpr const char *kPtxasRegisterUsageLevel =
    "tl.ptxas_register_usage_level";
45
46
static constexpr const char *kEnablePTXASVerboseOutput =
    "tl.enable_ptxas_verbose_output";
47
static constexpr const char *kDisableShuffleElect = "tl.disable_shuffle_elect";
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
/*!
 * \brief Whether to disable dynamic tail split
 *
 * kDisableDynamicTailSplit = "tl.disable_dynamic_tail_split"
 *
 */
static constexpr const char *kDisableDynamicTailSplit =
    "tl.disable_dynamic_tail_split";

/*!
 * \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";

69
70
71
72
73
74
75
76
/*!
 * \brief Get the type of the CUDA tensor map
 *
 * DataType cuTensorMapType()
 *
 */
DataType cuTensorMapType();

77
78
79
/*!
 * \brief tvm intrinsics for TMADescriptor creation for tiled load
 *
80
 * CuTensorMap* create_tma_descriptor(data_type, rank, global_addr,
81
82
 * global_shape..., global_stride..., smem_box..., smem_stride..., interleave,
 * swizzle, l2_promotion, oob_fill)
83
84
 *
 */
85
TVM_DLL const Op &create_tma_descriptor();
86
87
88
89

/*!
 * \brief tvm intrinsics for TMADescriptor creation for image to column load
 *
90
 * CuTensorMap* create_tma_im2col_descriptor(data_type, rank, global_addr,
91
92
93
 * global_shape..., global_stride..., elem_stride..., lower_corner...,
 * upper_corner..., smme_box_pixel, smem_box_channel, interleave, swizzle,
 * l2_promotion, oob_fill)
94
95
 *
 */
96
TVM_DLL const Op &create_tma_im2col_descriptor();
97
98
99
100

/*!
 * \brief Create a list of mbarrier with num_threads
 *
101
 * create_list_of_mbarrier(num_threads0, num_threads1, ...)
102
103
 *
 */
104
TVM_DLL const Op &create_list_of_mbarrier();
105
106
107
108
109
110
111

/*!
 * \brief Get the mbarrier with barrier_id
 *
 * int64_t* GetMBarrier(barrier_id)
 *
 */
112
TVM_DLL const Op &get_mbarrier();
113
114

/*!
115
116
 * \brief tvm intrinsics for loading data from global tensor descriptor to
 * shared memory
117
 *
118
 * tma_load(descriptor, mbarrier, smem_data, coord_0, coord_1, ...)
119
120
 *
 */
121
TVM_DLL const Op &tma_load();
122
123

/*!
124
125
 * \brief tvm intrinsics for loading image from global tensor to columns in
 * shared memory
126
 *
127
 * tma_load(descriptor, mbarrier, smem_data, coord_0, coord_1, ...,
128
 * image_offset, ...)
129
130
 *
 */
131
TVM_DLL const Op &tma_load_im2col();
132
133

/*!
134
135
 * \brief tvm intrinsics for storing data from shared memory to global tensor
 * descriptor
136
 *
137
 * tma_store(descriptor, smem_data, coord_0, coord_1, ...)
138
139
 *
 */
140
TVM_DLL const Op &tma_store();
141

142
143
144
145
146
147
148
149
/*!
 * \brief tvm intrinsics for barrier initialization fence
 *
 * ptx_fence_barrier_init()
 *
 */
const Op &ptx_fence_barrier_init();

150
151
152
/*!
 * \brief tvm intrinsics for mbarrier wait with parity bit
 *
153
 * mbarrier_wait_parity(mbarrier, parity)
154
155
 *
 */
156
TVM_DLL const Op &mbarrier_wait_parity();
157
158
159
160

/*!
 * \brief tvm intrinsics for mbarrier expect tx
 *
161
 * mbarrier_expect_tx(mbarrier, transaction_bytes)
162
163
 *
 */
164
TVM_DLL const Op &mbarrier_expect_tx();
165
166
167
168

/*!
 * \brief tvm intrinsics for ldmatrix
 *
169
 * ptx_ldmatrix(transposed, num, shared_addr, local_addr)
170
171
 *
 */
172
TVM_DLL const Op &ptx_ldmatrix();
173
174
175
176

/*!
 * \brief tvm intrinsics for stmatrix
 *
177
 * ptx_ldmatrix(transposed, num, shared_addr, int32_values...)
178
179
 *
 */
180
TVM_DLL const Op &ptx_stmatrix();
181

182
183
184
185
186
187
188
189
/*!
 * \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();

190
191
192
/*!
 * \brief Pack two b16 value into a b32 value
 *
193
 * int32 pack_b16(b16_value, b16_value)
194
195
 *
 */
196
TVM_DLL const Op &pack_b16();
197
198
199
200
201
202
203

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

206
207
208
/*!
 * \brief Indicate arrival of warp issuing TMA_STORE
 *
209
 * tma_store_arrive()
210
211
 *
 */
212
TVM_DLL const Op &tma_store_arrive();
213
214
215
216

/*!
 * \brief Wait for TMA_STORE to finish
 *
217
 * tma_store_wait()
218
219
 *
 */
220
TVM_DLL const Op &tma_store_wait();
221

222
223
224
225
226
227
/*!
 * \brief Set reg hint for warp-specialized branched
 *
 * SetMaxNRegInc(num_reg, is_inc)
 *
 */
228
TVM_DLL const Op &set_max_nreg();
229

230
231
232
/*!
 * \brief No set reg hint for warp-specialized branched
 *
233
 * no_set_max_nreg()
234
235
 *
 */
236
TVM_DLL const Op &no_set_max_nreg();
237

238
239
240
/*!
 * \brief Wait the previous wgmma to finish
 *
241
 * wait_wgmma(num_mma)
242
243
 *
 */
244
TVM_DLL const Op &wait_wgmma();
245

246
247
248
249
250
251
/*!
 * \brief Synchronize all threads in a grid
 *
 * sync_grid()
 *
 */
252
TVM_DLL const Op &sync_grid();
253
254
255
256
257
258
259

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

262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
/*!
 * \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();

310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
/*!
 * \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();

325
326
327
328
329
330
331
/*!
 * \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();

332
333
} // namespace tl
} // namespace tvm
334

335
#endif //  TVM_TL_OP_BUILTIN_H_