1/*
2 * Licensed to the Apache Software Foundation (ASF) under one
3 * or more contributor license agreements. See the NOTICE file
4 * distributed with this work for additional information
5 * regarding copyright ownership. The ASF licenses this file
6 * to you under the Apache License, Version 2.0 (the
7 * "License"); you may not use this file except in compliance
8 * with the License. You may obtain a copy of the License at
9 *
10 * http://www.apache.org/licenses/LICENSE-2.0
11 *
12 * Unless required by applicable law or agreed to in writing,
13 * software distributed under the License is distributed on an
14 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15 * KIND, either express or implied. See the License for the
16 * specific language governing permissions and limitations
17 * under the License.
18 */
19
20/*!
21 * \file tir/op/builtin.cc
22 *
23 * builtin intrinsic operators.
24 */
25#include <tvm/runtime/registry.h>
26#include <tvm/tir/builtin.h>
27#include <tvm/tir/op.h>
28#include <tvm/tir/op_attr_types.h>
29
30namespace tvm {
31namespace tir {
32namespace builtin {
33
34#define TIR_DEFINE_BUILTIN_FUNC(OpName) \
35 const Op& OpName() { \
36 static const Op& op = Op::Get("tir." #OpName); \
37 return op; \
38 } \
39 TVM_TIR_REGISTER_OP(#OpName)
40
41TIR_DEFINE_BUILTIN_FUNC(reinterpret)
42 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure))
43 .set_num_inputs(1);
44
45TIR_DEFINE_BUILTIN_FUNC(ret)
46 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kControlJump))
47 .set_num_inputs(1);
48
49TIR_DEFINE_BUILTIN_FUNC(likely)
50 .set_num_inputs(1)
51 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kExprAnnotation))
52 .set_attr<TVectorizable>("TVectorizable", true);
53
54TIR_DEFINE_BUILTIN_FUNC(bitwise_and)
55 .set_num_inputs(2)
56 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure))
57 .set_attr<TVectorizable>("TVectorizable", true);
58
59TIR_DEFINE_BUILTIN_FUNC(bitwise_or)
60 .set_num_inputs(2)
61 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure))
62 .set_attr<TVectorizable>("TVectorizable", true);
63
64TIR_DEFINE_BUILTIN_FUNC(bitwise_xor)
65 .set_num_inputs(2)
66 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure))
67 .set_attr<TVectorizable>("TVectorizable", true);
68
69TIR_DEFINE_BUILTIN_FUNC(bitwise_not)
70 .set_num_inputs(1)
71 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure))
72 .set_attr<TVectorizable>("TVectorizable", true);
73
74TIR_DEFINE_BUILTIN_FUNC(shift_left)
75 .set_num_inputs(2)
76 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure))
77 .set_attr<TVectorizable>("TVectorizable", true);
78
79TIR_DEFINE_BUILTIN_FUNC(shift_right)
80 .set_num_inputs(2)
81 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure))
82 .set_attr<TVectorizable>("TVectorizable", true);
83
84TIR_DEFINE_BUILTIN_FUNC(large_uint_imm)
85 .set_num_inputs(2)
86 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure));
87
88TIR_DEFINE_BUILTIN_FUNC(address_of)
89 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure))
90 .set_num_inputs(1);
91
92TIR_DEFINE_BUILTIN_FUNC(if_then_else)
93 .set_num_inputs(3)
94 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure));
95
96TIR_DEFINE_BUILTIN_FUNC(q_multiply_shift)
97 .set_num_inputs(3)
98 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure))
99 .set_attr<TVectorizable>("TVectorizable", true);
100
101TIR_DEFINE_BUILTIN_FUNC(q_multiply_shift_per_axis)
102 .set_num_inputs(7)
103 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure))
104 .set_attr<TVectorizable>("TVectorizable", true);
105
106TIR_DEFINE_BUILTIN_FUNC(isnullptr).set_num_inputs(1).set_attr<TCallEffectKind>(
107 "TCallEffectKind", Integer(CallEffectKind::kPure));
108
109TIR_DEFINE_BUILTIN_FUNC(isnan).set_num_inputs(1).set_attr<TCallEffectKind>(
110 "TCallEffectKind", Integer(CallEffectKind::kPure));
111
112TIR_DEFINE_BUILTIN_FUNC(popcount)
113 .set_num_inputs(1)
114 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure))
115 .set_attr<TVectorizable>("TVectorizable", true);
116
117TIR_DEFINE_BUILTIN_FUNC(fma)
118 .set_num_inputs(3)
119 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure))
120 .set_attr<TVectorizable>("TVectorizable", true);
121
122TIR_DEFINE_BUILTIN_FUNC(call_extern)
123 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
124
125TIR_DEFINE_BUILTIN_FUNC(call_pure_extern)
126 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure));
127
128TIR_DEFINE_BUILTIN_FUNC(call_llvm_intrin)
129 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
130
131TIR_DEFINE_BUILTIN_FUNC(call_llvm_pure_intrin)
132 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure));
133
134TIR_DEFINE_BUILTIN_FUNC(call_spirv_pure_glsl450)
135 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure));
136
137TIR_DEFINE_BUILTIN_FUNC(prefetch).set_attr<TCallEffectKind>("TCallEffectKind",
138 Integer(CallEffectKind::kOpaque));
139
140TIR_DEFINE_BUILTIN_FUNC(tvm_access_ptr)
141 .set_num_inputs(5)
142 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kSpecialCallArg));
143
144TIR_DEFINE_BUILTIN_FUNC(tvm_static_handle)
145 .set_num_inputs(0)
146 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kSpecialCallArg));
147
148TIR_DEFINE_BUILTIN_FUNC(tvm_context_id)
149 .set_num_inputs(0)
150 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kReadState));
151
152TIR_DEFINE_BUILTIN_FUNC(tvm_tuple).set_attr<TCallEffectKind>("TCallEffectKind",
153 Integer(CallEffectKind::kEmbedInfo));
154
155TIR_DEFINE_BUILTIN_FUNC(tvm_struct_get)
156 .set_num_inputs(3)
157 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kReadState));
158
159TIR_DEFINE_BUILTIN_FUNC(tvm_struct_set)
160 .set_num_inputs(4)
161 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kUpdateState));
162
163TIR_DEFINE_BUILTIN_FUNC(lookup_param)
164 .set_num_inputs(4)
165 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kUpdateState));
166
167TIR_DEFINE_BUILTIN_FUNC(tvm_throw_last_error)
168 .set_num_inputs(0)
169 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
170
171TIR_DEFINE_BUILTIN_FUNC(tvm_stack_alloca)
172 .set_num_inputs(2)
173 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
174
175TIR_DEFINE_BUILTIN_FUNC(tvm_stack_make_shape)
176 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
177
178TIR_DEFINE_BUILTIN_FUNC(tvm_stack_make_array)
179 .set_num_inputs(6)
180 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
181
182// When num_inputs are not set, the function is assumed to be variable length.
183TIR_DEFINE_BUILTIN_FUNC(tvm_call_packed)
184 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque))
185 .set_attr<TScriptPrinterName>("TScriptPrinterName", String("call_packed"), /*plevel=*/20);
186
187TIR_DEFINE_BUILTIN_FUNC(tvm_call_cpacked)
188 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque))
189 .set_attr<TScriptPrinterName>("TScriptPrinterName", String("call_cpacked"), /*plevel=*/20);
190
191TIR_DEFINE_BUILTIN_FUNC(tvm_call_trace_packed)
192 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
193
194TIR_DEFINE_BUILTIN_FUNC(tvm_check_return)
195 .set_num_inputs(3)
196 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure));
197
198TIR_DEFINE_BUILTIN_FUNC(tvm_thread_context)
199 .set_num_inputs(1)
200 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
201
202TIR_DEFINE_BUILTIN_FUNC(tvm_call_packed_lowered)
203 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque))
204 .set_attr<TScriptPrinterName>("TScriptPrinterName", String("call_packed_lowered"),
205 /*plevel=*/20);
206
207TIR_DEFINE_BUILTIN_FUNC(tvm_call_cpacked_lowered)
208 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque))
209 .set_attr<TScriptPrinterName>("TScriptPrinterName", String("call_cpacked_lowered"),
210 /*plevel=*/20);
211
212TIR_DEFINE_BUILTIN_FUNC(tvm_call_trace_packed_lowered)
213 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
214
215// TODO(tvm-team) revisit storage sync once we have a good memory hierachy structure.
216TIR_DEFINE_BUILTIN_FUNC(tvm_storage_sync)
217 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
218
219TIR_DEFINE_BUILTIN_FUNC(tvm_warp_shuffle)
220 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
221
222TIR_DEFINE_BUILTIN_FUNC(tvm_warp_shuffle_up)
223 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
224
225TIR_DEFINE_BUILTIN_FUNC(tvm_warp_shuffle_down)
226 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
227
228TIR_DEFINE_BUILTIN_FUNC(tvm_warp_activemask)
229 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
230
231TIR_DEFINE_BUILTIN_FUNC(tvm_global_barrier_kinit)
232 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
233
234TIR_DEFINE_BUILTIN_FUNC(tvm_thread_allreduce)
235 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
236
237TIR_DEFINE_BUILTIN_FUNC(tvm_load_matrix_sync)
238 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kReadState));
239
240TIR_DEFINE_BUILTIN_FUNC(tvm_mma_sync)
241 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
242
243TIR_DEFINE_BUILTIN_FUNC(tvm_bmma_sync)
244 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
245
246TIR_DEFINE_BUILTIN_FUNC(tvm_fill_fragment)
247 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
248
249TIR_DEFINE_BUILTIN_FUNC(tvm_store_matrix_sync)
250 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
251
252TIR_DEFINE_BUILTIN_FUNC(ptx_mma).set_attr<TCallEffectKind>("TCallEffectKind",
253 Integer(CallEffectKind::kOpaque));
254
255TIR_DEFINE_BUILTIN_FUNC(ptx_mma_sp)
256 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
257
258TIR_DEFINE_BUILTIN_FUNC(ptx_ldmatrix)
259 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
260
261TIR_DEFINE_BUILTIN_FUNC(ptx_cp_async)
262 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
263
264TIR_DEFINE_BUILTIN_FUNC(ptx_commit_group)
265 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
266
267TIR_DEFINE_BUILTIN_FUNC(ptx_wait_group)
268 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
269
270TIR_DEFINE_BUILTIN_FUNC(mma_store).set_attr<TCallEffectKind>("TCallEffectKind",
271 Integer(CallEffectKind::kOpaque));
272
273TIR_DEFINE_BUILTIN_FUNC(mma_fill).set_attr<TCallEffectKind>("TCallEffectKind",
274 Integer(CallEffectKind::kOpaque));
275
276TIR_DEFINE_BUILTIN_FUNC(vectorhigh)
277 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure));
278
279TIR_DEFINE_BUILTIN_FUNC(vectorlow).set_attr<TCallEffectKind>("TCallEffectKind",
280 Integer(CallEffectKind::kPure));
281
282TIR_DEFINE_BUILTIN_FUNC(vectorcombine)
283 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure));
284
285TIR_DEFINE_BUILTIN_FUNC(atomic_add)
286 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
287
288TIR_DEFINE_BUILTIN_FUNC(nd_mem_alloc_with_scope)
289 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
290
291TIR_DEFINE_BUILTIN_FUNC(texture2d_store)
292 .set_attr<TVectorizable>("TVectorizable", true)
293 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
294
295TIR_DEFINE_BUILTIN_FUNC(texture2d_load)
296 .set_attr<TVectorizable>("TVectorizable", true)
297 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
298
299TIR_DEFINE_BUILTIN_FUNC(dma_copy).set_attr<TCallEffectKind>("TCallEffectKind",
300 Integer(CallEffectKind::kOpaque));
301
302TIR_DEFINE_BUILTIN_FUNC(dma_wait).set_attr<TCallEffectKind>("TCallEffectKind",
303 Integer(CallEffectKind::kOpaque));
304
305TIR_DEFINE_BUILTIN_FUNC(assume)
306 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kEmbedInfo))
307 .set_num_inputs(1);
308
309TIR_DEFINE_BUILTIN_FUNC(undef)
310 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kReadState))
311 .set_num_inputs(0);
312
313TIR_DEFINE_BUILTIN_FUNC(start_profile_intrinsic)
314 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure));
315
316TIR_DEFINE_BUILTIN_FUNC(end_profile_intrinsic)
317 .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure));
318
319} // namespace builtin
320} // namespace tir
321} // namespace tvm
322