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 | |
30 | namespace tvm { |
31 | namespace tir { |
32 | namespace 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 | |
41 | TIR_DEFINE_BUILTIN_FUNC(reinterpret) |
42 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kPure)) |
43 | .set_num_inputs(1); |
44 | |
45 | TIR_DEFINE_BUILTIN_FUNC(ret) |
46 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kControlJump)) |
47 | .set_num_inputs(1); |
48 | |
49 | TIR_DEFINE_BUILTIN_FUNC(likely) |
50 | .set_num_inputs(1) |
51 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kExprAnnotation)) |
52 | .set_attr<TVectorizable>("TVectorizable" , true); |
53 | |
54 | TIR_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 | |
59 | TIR_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 | |
64 | TIR_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 | |
69 | TIR_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 | |
74 | TIR_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 | |
79 | TIR_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 | |
84 | TIR_DEFINE_BUILTIN_FUNC(large_uint_imm) |
85 | .set_num_inputs(2) |
86 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kPure)); |
87 | |
88 | TIR_DEFINE_BUILTIN_FUNC(address_of) |
89 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kPure)) |
90 | .set_num_inputs(1); |
91 | |
92 | TIR_DEFINE_BUILTIN_FUNC(if_then_else) |
93 | .set_num_inputs(3) |
94 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kPure)); |
95 | |
96 | TIR_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 | |
101 | TIR_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 | |
106 | TIR_DEFINE_BUILTIN_FUNC(isnullptr).set_num_inputs(1).set_attr<TCallEffectKind>( |
107 | "TCallEffectKind" , Integer(CallEffectKind::kPure)); |
108 | |
109 | TIR_DEFINE_BUILTIN_FUNC(isnan).set_num_inputs(1).set_attr<TCallEffectKind>( |
110 | "TCallEffectKind" , Integer(CallEffectKind::kPure)); |
111 | |
112 | TIR_DEFINE_BUILTIN_FUNC(popcount) |
113 | .set_num_inputs(1) |
114 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kPure)) |
115 | .set_attr<TVectorizable>("TVectorizable" , true); |
116 | |
117 | TIR_DEFINE_BUILTIN_FUNC(fma) |
118 | .set_num_inputs(3) |
119 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kPure)) |
120 | .set_attr<TVectorizable>("TVectorizable" , true); |
121 | |
122 | TIR_DEFINE_BUILTIN_FUNC(call_extern) |
123 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kOpaque)); |
124 | |
125 | TIR_DEFINE_BUILTIN_FUNC(call_pure_extern) |
126 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kPure)); |
127 | |
128 | TIR_DEFINE_BUILTIN_FUNC(call_llvm_intrin) |
129 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kOpaque)); |
130 | |
131 | TIR_DEFINE_BUILTIN_FUNC(call_llvm_pure_intrin) |
132 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kPure)); |
133 | |
134 | TIR_DEFINE_BUILTIN_FUNC(call_spirv_pure_glsl450) |
135 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kPure)); |
136 | |
137 | TIR_DEFINE_BUILTIN_FUNC(prefetch).set_attr<TCallEffectKind>("TCallEffectKind" , |
138 | Integer(CallEffectKind::kOpaque)); |
139 | |
140 | TIR_DEFINE_BUILTIN_FUNC(tvm_access_ptr) |
141 | .set_num_inputs(5) |
142 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kSpecialCallArg)); |
143 | |
144 | TIR_DEFINE_BUILTIN_FUNC(tvm_static_handle) |
145 | .set_num_inputs(0) |
146 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kSpecialCallArg)); |
147 | |
148 | TIR_DEFINE_BUILTIN_FUNC(tvm_context_id) |
149 | .set_num_inputs(0) |
150 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kReadState)); |
151 | |
152 | TIR_DEFINE_BUILTIN_FUNC(tvm_tuple).set_attr<TCallEffectKind>("TCallEffectKind" , |
153 | Integer(CallEffectKind::kEmbedInfo)); |
154 | |
155 | TIR_DEFINE_BUILTIN_FUNC(tvm_struct_get) |
156 | .set_num_inputs(3) |
157 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kReadState)); |
158 | |
159 | TIR_DEFINE_BUILTIN_FUNC(tvm_struct_set) |
160 | .set_num_inputs(4) |
161 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kUpdateState)); |
162 | |
163 | TIR_DEFINE_BUILTIN_FUNC(lookup_param) |
164 | .set_num_inputs(4) |
165 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kUpdateState)); |
166 | |
167 | TIR_DEFINE_BUILTIN_FUNC(tvm_throw_last_error) |
168 | .set_num_inputs(0) |
169 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kOpaque)); |
170 | |
171 | TIR_DEFINE_BUILTIN_FUNC(tvm_stack_alloca) |
172 | .set_num_inputs(2) |
173 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kOpaque)); |
174 | |
175 | TIR_DEFINE_BUILTIN_FUNC(tvm_stack_make_shape) |
176 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kOpaque)); |
177 | |
178 | TIR_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. |
183 | TIR_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 | |
187 | TIR_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 | |
191 | TIR_DEFINE_BUILTIN_FUNC(tvm_call_trace_packed) |
192 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kOpaque)); |
193 | |
194 | TIR_DEFINE_BUILTIN_FUNC(tvm_check_return) |
195 | .set_num_inputs(3) |
196 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kPure)); |
197 | |
198 | TIR_DEFINE_BUILTIN_FUNC(tvm_thread_context) |
199 | .set_num_inputs(1) |
200 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kOpaque)); |
201 | |
202 | TIR_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 | |
207 | TIR_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 | |
212 | TIR_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. |
216 | TIR_DEFINE_BUILTIN_FUNC(tvm_storage_sync) |
217 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kOpaque)); |
218 | |
219 | TIR_DEFINE_BUILTIN_FUNC(tvm_warp_shuffle) |
220 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kOpaque)); |
221 | |
222 | TIR_DEFINE_BUILTIN_FUNC(tvm_warp_shuffle_up) |
223 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kOpaque)); |
224 | |
225 | TIR_DEFINE_BUILTIN_FUNC(tvm_warp_shuffle_down) |
226 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kOpaque)); |
227 | |
228 | TIR_DEFINE_BUILTIN_FUNC(tvm_warp_activemask) |
229 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kOpaque)); |
230 | |
231 | TIR_DEFINE_BUILTIN_FUNC(tvm_global_barrier_kinit) |
232 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kOpaque)); |
233 | |
234 | TIR_DEFINE_BUILTIN_FUNC(tvm_thread_allreduce) |
235 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kOpaque)); |
236 | |
237 | TIR_DEFINE_BUILTIN_FUNC(tvm_load_matrix_sync) |
238 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kReadState)); |
239 | |
240 | TIR_DEFINE_BUILTIN_FUNC(tvm_mma_sync) |
241 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kOpaque)); |
242 | |
243 | TIR_DEFINE_BUILTIN_FUNC(tvm_bmma_sync) |
244 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kOpaque)); |
245 | |
246 | TIR_DEFINE_BUILTIN_FUNC(tvm_fill_fragment) |
247 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kOpaque)); |
248 | |
249 | TIR_DEFINE_BUILTIN_FUNC(tvm_store_matrix_sync) |
250 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kOpaque)); |
251 | |
252 | TIR_DEFINE_BUILTIN_FUNC(ptx_mma).set_attr<TCallEffectKind>("TCallEffectKind" , |
253 | Integer(CallEffectKind::kOpaque)); |
254 | |
255 | TIR_DEFINE_BUILTIN_FUNC(ptx_mma_sp) |
256 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kOpaque)); |
257 | |
258 | TIR_DEFINE_BUILTIN_FUNC(ptx_ldmatrix) |
259 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kOpaque)); |
260 | |
261 | TIR_DEFINE_BUILTIN_FUNC(ptx_cp_async) |
262 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kOpaque)); |
263 | |
264 | TIR_DEFINE_BUILTIN_FUNC(ptx_commit_group) |
265 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kOpaque)); |
266 | |
267 | TIR_DEFINE_BUILTIN_FUNC(ptx_wait_group) |
268 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kOpaque)); |
269 | |
270 | TIR_DEFINE_BUILTIN_FUNC(mma_store).set_attr<TCallEffectKind>("TCallEffectKind" , |
271 | Integer(CallEffectKind::kOpaque)); |
272 | |
273 | TIR_DEFINE_BUILTIN_FUNC(mma_fill).set_attr<TCallEffectKind>("TCallEffectKind" , |
274 | Integer(CallEffectKind::kOpaque)); |
275 | |
276 | TIR_DEFINE_BUILTIN_FUNC(vectorhigh) |
277 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kPure)); |
278 | |
279 | TIR_DEFINE_BUILTIN_FUNC(vectorlow).set_attr<TCallEffectKind>("TCallEffectKind" , |
280 | Integer(CallEffectKind::kPure)); |
281 | |
282 | TIR_DEFINE_BUILTIN_FUNC(vectorcombine) |
283 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kPure)); |
284 | |
285 | TIR_DEFINE_BUILTIN_FUNC(atomic_add) |
286 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kOpaque)); |
287 | |
288 | TIR_DEFINE_BUILTIN_FUNC(nd_mem_alloc_with_scope) |
289 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kOpaque)); |
290 | |
291 | TIR_DEFINE_BUILTIN_FUNC(texture2d_store) |
292 | .set_attr<TVectorizable>("TVectorizable" , true) |
293 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kOpaque)); |
294 | |
295 | TIR_DEFINE_BUILTIN_FUNC(texture2d_load) |
296 | .set_attr<TVectorizable>("TVectorizable" , true) |
297 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kOpaque)); |
298 | |
299 | TIR_DEFINE_BUILTIN_FUNC(dma_copy).set_attr<TCallEffectKind>("TCallEffectKind" , |
300 | Integer(CallEffectKind::kOpaque)); |
301 | |
302 | TIR_DEFINE_BUILTIN_FUNC(dma_wait).set_attr<TCallEffectKind>("TCallEffectKind" , |
303 | Integer(CallEffectKind::kOpaque)); |
304 | |
305 | TIR_DEFINE_BUILTIN_FUNC(assume) |
306 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kEmbedInfo)) |
307 | .set_num_inputs(1); |
308 | |
309 | TIR_DEFINE_BUILTIN_FUNC(undef) |
310 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kReadState)) |
311 | .set_num_inputs(0); |
312 | |
313 | TIR_DEFINE_BUILTIN_FUNC(start_profile_intrinsic) |
314 | .set_attr<TCallEffectKind>("TCallEffectKind" , Integer(CallEffectKind::kPure)); |
315 | |
316 | TIR_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 | |