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

#ifndef TVM_TL_OP_BUILTIN_H_
#define TVM_TL_OP_BUILTIN_H_

#include "op.h"
11
#include <tvm/ir/transform.h>
12
13
14

namespace tvm {
namespace tl {
15
16
17
18
19

namespace attr {
static constexpr const char *kPaddingMap = "padding_map";
} // namespace attr

20
21
static constexpr const char *kDebugMergeSharedMemoryAllocations =
    "tl.debug_merge_shared_memory_allocations";
22
static constexpr const char *kDisableTMALower = "tl.disable_tma_lower";
23
24
static constexpr const char *kDisableSafeMemoryLegalize =
    "tl.disable_safe_memory_legalize";
25
26
static constexpr const char *kDisableWarpSpecialized =
    "tl.disable_warp_specialized";
27
28
static constexpr const char *kConfigIndexBitwidth = "tl.config_index_bitwidth";

29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
/*!
 * \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";

50
51
52
/*!
 * \brief tvm intrinsics for TMADescriptor creation for tiled load
 *
53
 * CuTensorMap* create_tma_descriptor(data_type, rank, global_addr,
54
55
 * global_shape..., global_stride..., smem_box..., smem_stride..., interleave,
 * swizzle, l2_promotion, oob_fill)
56
57
 *
 */
58
const Op &create_tma_descriptor();
59
60
61
62

/*!
 * \brief tvm intrinsics for TMADescriptor creation for image to column load
 *
63
 * CuTensorMap* create_tma_im2col_descriptor(data_type, rank, global_addr,
64
65
66
 * global_shape..., global_stride..., elem_stride..., lower_corner...,
 * upper_corner..., smme_box_pixel, smem_box_channel, interleave, swizzle,
 * l2_promotion, oob_fill)
67
68
 *
 */
69
const Op &create_tma_im2col_descriptor();
70
71
72
73

/*!
 * \brief Create a list of mbarrier with num_threads
 *
74
 * create_list_of_mbarrier(num_threads0, num_threads1, ...)
75
76
 *
 */
77
const Op &create_list_of_mbarrier();
78
79
80
81
82
83
84

/*!
 * \brief Get the mbarrier with barrier_id
 *
 * int64_t* GetMBarrier(barrier_id)
 *
 */
85
const Op &get_mbarrier();
86
87

/*!
88
89
 * \brief tvm intrinsics for loading data from global tensor descriptor to
 * shared memory
90
 *
91
 * tma_load(descriptor, mbarrier, smem_data, coord_0, coord_1, ...)
92
93
 *
 */
94
const Op &tma_load();
95
96

/*!
97
98
 * \brief tvm intrinsics for loading image from global tensor to columns in
 * shared memory
99
 *
100
 * tma_load(descriptor, mbarrier, smem_data, coord_0, coord_1, ...,
101
 * image_offset, ...)
102
103
 *
 */
104
const Op &tma_load_im2col();
105
106

/*!
107
108
 * \brief tvm intrinsics for storing data from shared memory to global tensor
 * descriptor
109
 *
110
 * tma_store(descriptor, smem_data, coord_0, coord_1, ...)
111
112
 *
 */
113
const Op &tma_store();
114
115
116
117

/*!
 * \brief tvm intrinsics for mbarrier wait with parity bit
 *
118
 * mbarrier_wait_parity(mbarrier, parity)
119
120
 *
 */
121
const Op &mbarrier_wait_parity();
122
123
124
125

/*!
 * \brief tvm intrinsics for mbarrier expect tx
 *
126
 * mbarrier_expect_tx(mbarrier, transaction_bytes)
127
128
 *
 */
129
const Op &mbarrier_expect_tx();
130
131
132
133

/*!
 * \brief tvm intrinsics for ldmatrix
 *
134
 * ptx_ldmatirx(transposed, num, shared_addr, local_addr)
135
136
 *
 */
137
const Op &ptx_ldmatirx();
138
139
140
141

/*!
 * \brief tvm intrinsics for stmatrix
 *
142
 * ptx_ldmatirx(transposed, num, shared_addr, int32_values...)
143
144
 *
 */
145
const Op &ptx_stmatirx();
146
147
148
149

/*!
 * \brief Pack two b16 value into a b32 value
 *
150
 * int32 pack_b16(b16_value, b16_value)
151
152
 *
 */
153
const Op &pack_b16();
154
155
156
157

/*!
 * \brief Similar to __syncthreads(), but can be used to sync partial threads
 *
158
 * sync_thread_partial(num_partial_threads or mbarrier)
159
160
 *
 */
161
const Op &sync_thread_partial();
162
163
164
165
166
167
168

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

171
172
173
/*!
 * \brief Indicate arrival of warp issuing TMA_STORE
 *
174
 * tma_store_arrive()
175
176
 *
 */
177
const Op &tma_store_arrive();
178
179
180
181

/*!
 * \brief Wait for TMA_STORE to finish
 *
182
 * tma_store_wait()
183
184
 *
 */
185
const Op &tma_store_wait();
186

187
188
189
190
191
192
/*!
 * \brief Set reg hint for warp-specialized branched
 *
 * SetMaxNRegInc(num_reg, is_inc)
 *
 */
193
const Op &set_max_nreg();
194

195
196
197
/*!
 * \brief No set reg hint for warp-specialized branched
 *
198
 * no_set_max_nreg()
199
200
 *
 */
201
const Op &no_set_max_nreg();
202

203
204
205
/*!
 * \brief Wait the previous wgmma to finish
 *
206
 * wait_wgmma(num_mma)
207
208
 *
 */
209
const Op &wait_wgmma();
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258

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

259
260
} // namespace tl
} // namespace tvm
261

262
#endif //  TVM_TL_OP_BUILTIN_H_