builtin.h 7.74 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
14

namespace tvm {
namespace tl {
15
16
17

namespace attr {
static constexpr const char *kPaddingMap = "padding_map";
18
19
static constexpr const char *kWarpSpecializationScope =
    "kWarpSpecializationScope";
20
21
} // namespace attr

22
23
static constexpr const char *kDebugMergeSharedMemoryAllocations =
    "tl.debug_merge_shared_memory_allocations";
24
static constexpr const char *kDisableTMALower = "tl.disable_tma_lower";
25
26
static constexpr const char *kDisableSafeMemoryLegalize =
    "tl.disable_safe_memory_legalize";
27
28
static constexpr const char *kDisableWarpSpecialized =
    "tl.disable_warp_specialized";
29
static constexpr const char *kConfigIndexBitwidth = "tl.config_index_bitwidth";
30
31
static constexpr const char *kEnableAggressiveSharedMemoryMerge =
    "tl.enable_aggressive_shared_memory_merge";
32
static constexpr const char *kDisableFastMath = "tl.disable_fast_math";
33
34
static constexpr const char *kPtxasRegisterUsageLevel =
    "tl.ptxas_register_usage_level";
35
36
static constexpr const char *kEnablePTXASVerboseOutput =
    "tl.enable_ptxas_verbose_output";
37
static constexpr const char *kDisableShuffleElect = "tl.disable_shuffle_elect";
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
/*!
 * \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";

59
60
61
62
63
64
65
66
/*!
 * \brief Get the type of the CUDA tensor map
 *
 * DataType cuTensorMapType()
 *
 */
DataType cuTensorMapType();

67
68
69
/*!
 * \brief tvm intrinsics for TMADescriptor creation for tiled load
 *
70
 * CuTensorMap* create_tma_descriptor(data_type, rank, global_addr,
71
72
 * global_shape..., global_stride..., smem_box..., smem_stride..., interleave,
 * swizzle, l2_promotion, oob_fill)
73
74
 *
 */
75
TVM_DLL const Op &create_tma_descriptor();
76
77
78
79

/*!
 * \brief tvm intrinsics for TMADescriptor creation for image to column load
 *
80
 * CuTensorMap* create_tma_im2col_descriptor(data_type, rank, global_addr,
81
82
83
 * global_shape..., global_stride..., elem_stride..., lower_corner...,
 * upper_corner..., smme_box_pixel, smem_box_channel, interleave, swizzle,
 * l2_promotion, oob_fill)
84
85
 *
 */
86
TVM_DLL const Op &create_tma_im2col_descriptor();
87
88
89
90

/*!
 * \brief Create a list of mbarrier with num_threads
 *
91
 * create_list_of_mbarrier(num_threads0, num_threads1, ...)
92
93
 *
 */
94
TVM_DLL const Op &create_list_of_mbarrier();
95
96
97
98
99
100
101

/*!
 * \brief Get the mbarrier with barrier_id
 *
 * int64_t* GetMBarrier(barrier_id)
 *
 */
102
TVM_DLL const Op &get_mbarrier();
103
104

/*!
105
106
 * \brief tvm intrinsics for loading data from global tensor descriptor to
 * shared memory
107
 *
108
 * tma_load(descriptor, mbarrier, smem_data, coord_0, coord_1, ...)
109
110
 *
 */
111
TVM_DLL const Op &tma_load();
112
113

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

/*!
124
125
 * \brief tvm intrinsics for storing data from shared memory to global tensor
 * descriptor
126
 *
127
 * tma_store(descriptor, smem_data, coord_0, coord_1, ...)
128
129
 *
 */
130
TVM_DLL const Op &tma_store();
131

132
133
134
135
136
137
138
139
/*!
 * \brief tvm intrinsics for barrier initialization fence
 *
 * ptx_fence_barrier_init()
 *
 */
const Op &ptx_fence_barrier_init();

140
141
142
/*!
 * \brief tvm intrinsics for mbarrier wait with parity bit
 *
143
 * mbarrier_wait_parity(mbarrier, parity)
144
145
 *
 */
146
TVM_DLL const Op &mbarrier_wait_parity();
147
148
149
150

/*!
 * \brief tvm intrinsics for mbarrier expect tx
 *
151
 * mbarrier_expect_tx(mbarrier, transaction_bytes)
152
153
 *
 */
154
TVM_DLL const Op &mbarrier_expect_tx();
155
156
157
158

/*!
 * \brief tvm intrinsics for ldmatrix
 *
159
 * ptx_ldmatrix(transposed, num, shared_addr, local_addr)
160
161
 *
 */
162
TVM_DLL const Op &ptx_ldmatrix();
163
164
165
166

/*!
 * \brief tvm intrinsics for stmatrix
 *
167
 * ptx_ldmatrix(transposed, num, shared_addr, int32_values...)
168
169
 *
 */
170
TVM_DLL const Op &ptx_stmatrix();
171
172
173
174

/*!
 * \brief Pack two b16 value into a b32 value
 *
175
 * int32 pack_b16(b16_value, b16_value)
176
177
 *
 */
178
TVM_DLL const Op &pack_b16();
179
180
181
182
183
184
185

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

188
189
190
/*!
 * \brief Indicate arrival of warp issuing TMA_STORE
 *
191
 * tma_store_arrive()
192
193
 *
 */
194
TVM_DLL const Op &tma_store_arrive();
195
196
197
198

/*!
 * \brief Wait for TMA_STORE to finish
 *
199
 * tma_store_wait()
200
201
 *
 */
202
TVM_DLL const Op &tma_store_wait();
203

204
205
206
207
208
209
/*!
 * \brief Set reg hint for warp-specialized branched
 *
 * SetMaxNRegInc(num_reg, is_inc)
 *
 */
210
TVM_DLL const Op &set_max_nreg();
211

212
213
214
/*!
 * \brief No set reg hint for warp-specialized branched
 *
215
 * no_set_max_nreg()
216
217
 *
 */
218
TVM_DLL const Op &no_set_max_nreg();
219

220
221
222
/*!
 * \brief Wait the previous wgmma to finish
 *
223
 * wait_wgmma(num_mma)
224
225
 *
 */
226
TVM_DLL const Op &wait_wgmma();
227

228
229
230
231
232
233
/*!
 * \brief Synchronize all threads in a grid
 *
 * sync_grid()
 *
 */
234
TVM_DLL const Op &sync_grid();
235
236
237
238
239
240
241

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

244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
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
/*!
 * \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();

292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
/*!
 * \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();

307
308
309
310
311
312
313
/*!
 * \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();

314
315
} // namespace tl
} // namespace tvm
316

317
#endif //  TVM_TL_OP_BUILTIN_H_