1use crate::{
4 ir::{
5 attribute::{
6 DenseElementsAttribute, DenseI32ArrayAttribute, IntegerAttribute, StringAttribute,
7 },
8 operation::OperationBuilder,
9 r#type::RankedTensorType,
10 Block, Identifier, Location, Operation, Type, Value,
11 },
12 Context, Error,
13};
14
15pub fn assert<'c>(
17 context: &'c Context,
18 argument: Value<'c, '_>,
19 message: &str,
20 location: Location<'c>,
21) -> Operation<'c> {
22 OperationBuilder::new("cf.assert", location)
23 .add_attributes(&[(
24 Identifier::new(context, "msg"),
25 StringAttribute::new(context, message).into(),
26 )])
27 .add_operands(&[argument])
28 .build()
29 .expect("valid operation")
30}
31
32pub fn br<'c>(
34 successor: &Block<'c>,
35 destination_operands: &[Value<'c, '_>],
36 location: Location<'c>,
37) -> Operation<'c> {
38 OperationBuilder::new("cf.br", location)
39 .add_operands(destination_operands)
40 .add_successors(&[successor])
41 .build()
42 .expect("valid operation")
43}
44
45pub fn cond_br<'c>(
47 context: &'c Context,
48 condition: Value<'c, '_>,
49 true_successor: &Block<'c>,
50 false_successor: &Block<'c>,
51 true_successor_operands: &[Value<'c, '_>],
52 false_successor_operands: &[Value<'c, '_>],
53 location: Location<'c>,
54) -> Operation<'c> {
55 OperationBuilder::new("cf.cond_br", location)
56 .add_attributes(&[(
57 Identifier::new(context, "operand_segment_sizes"),
58 DenseI32ArrayAttribute::new(
59 context,
60 &[
61 1,
62 true_successor.argument_count() as i32,
63 false_successor.argument_count() as i32,
64 ],
65 )
66 .into(),
67 )])
68 .add_operands(
69 &[condition]
70 .into_iter()
71 .chain(true_successor_operands.iter().copied())
72 .chain(false_successor_operands.iter().copied())
73 .collect::<Vec<_>>(),
74 )
75 .add_successors(&[true_successor, false_successor])
76 .build()
77 .expect("valid operation")
78}
79
80pub fn switch<'c>(
82 context: &'c Context,
83 case_values: &[i64],
84 flag: Value<'c, '_>,
85 flag_type: Type<'c>,
86 default_destination: (&Block<'c>, &[Value<'c, '_>]),
87 case_destinations: &[(&Block<'c>, &[Value<'c, '_>])],
88 location: Location<'c>,
89) -> Result<Operation<'c>, Error> {
90 let (destinations, operands): (Vec<_>, Vec<_>) = [default_destination]
91 .into_iter()
92 .chain(case_destinations.iter().copied())
93 .unzip();
94
95 Ok(OperationBuilder::new("cf.switch", location)
96 .add_attributes(&[
97 (
98 Identifier::new(context, "case_values"),
99 DenseElementsAttribute::new(
100 RankedTensorType::new(&[case_values.len() as u64], flag_type, None).into(),
101 &case_values
102 .iter()
103 .map(|value| IntegerAttribute::new(flag_type, *value).into())
104 .collect::<Vec<_>>(),
105 )?
106 .into(),
107 ),
108 (
109 Identifier::new(context, "case_operand_segments"),
110 DenseI32ArrayAttribute::new(
111 context,
112 &case_destinations
113 .iter()
114 .map(|(_, operands)| operands.len() as i32)
115 .collect::<Vec<_>>(),
116 )
117 .into(),
118 ),
119 (
120 Identifier::new(context, "operand_segment_sizes"),
121 DenseI32ArrayAttribute::new(
122 context,
123 &[
124 1,
125 default_destination.1.len() as i32,
126 case_destinations
127 .iter()
128 .map(|(_, operands)| operands.len() as i32)
129 .sum(),
130 ],
131 )
132 .into(),
133 ),
134 ])
135 .add_operands(
136 &[flag]
137 .into_iter()
138 .chain(operands.into_iter().flatten().copied())
139 .collect::<Vec<_>>(),
140 )
141 .add_successors(&destinations)
142 .build()
143 .expect("valid operation"))
144}
145
146#[cfg(test)]
147mod tests {
148 use super::*;
149 use crate::{
150 dialect::{
151 arith::{self, CmpiPredicate},
152 func, index,
153 },
154 ir::{
155 attribute::{IntegerAttribute, StringAttribute, TypeAttribute},
156 r#type::{FunctionType, IntegerType, Type},
157 Block, Module, Region,
158 },
159 test::load_all_dialects,
160 Context,
161 };
162
163 #[test]
164 fn compile_assert() {
165 let context = Context::new();
166 load_all_dialects(&context);
167
168 let location = Location::unknown(&context);
169 let module = Module::new(location);
170 let bool_type: Type = IntegerType::new(&context, 1).into();
171
172 module.body().append_operation(func::func(
173 &context,
174 StringAttribute::new(&context, "foo"),
175 TypeAttribute::new(FunctionType::new(&context, &[], &[]).into()),
176 {
177 let block = Block::new(&[]);
178 let operand = block
179 .append_operation(arith::constant(
180 &context,
181 IntegerAttribute::new(bool_type, 1).into(),
182 location,
183 ))
184 .result(0)
185 .unwrap()
186 .into();
187
188 block.append_operation(assert(&context, operand, "assert message", location));
189
190 block.append_operation(func::r#return(&[], location));
191
192 let region = Region::new();
193 region.append_block(block);
194 region
195 },
196 &[],
197 location,
198 ));
199
200 assert!(module.as_operation().verify());
201 insta::assert_snapshot!(module.as_operation());
202 }
203
204 #[test]
205 fn compile_br() {
206 let context = Context::new();
207 load_all_dialects(&context);
208
209 let location = Location::unknown(&context);
210 let module = Module::new(location);
211 let index_type = Type::index(&context);
212
213 module.body().append_operation(func::func(
214 &context,
215 StringAttribute::new(&context, "foo"),
216 TypeAttribute::new(FunctionType::new(&context, &[], &[]).into()),
217 {
218 let block = Block::new(&[]);
219 let dest_block = Block::new(&[(index_type, location)]);
220 let operand = block
221 .append_operation(index::constant(
222 &context,
223 IntegerAttribute::new(index_type, 1),
224 location,
225 ))
226 .result(0)
227 .unwrap();
228
229 block.append_operation(br(&dest_block, &[operand.into()], location));
230
231 dest_block.append_operation(func::r#return(&[], location));
232
233 let region = Region::new();
234 region.append_block(block);
235 region.append_block(dest_block);
236 region
237 },
238 &[],
239 location,
240 ));
241
242 assert!(module.as_operation().verify());
243 insta::assert_snapshot!(module.as_operation());
244 }
245
246 #[test]
247 fn compile_cond_br() {
248 let context = Context::new();
249 load_all_dialects(&context);
250
251 let location = Location::unknown(&context);
252 let module = Module::new(location);
253 let index_type = Type::index(&context);
254
255 module.body().append_operation(func::func(
256 &context,
257 StringAttribute::new(&context, "foo"),
258 TypeAttribute::new(FunctionType::new(&context, &[], &[]).into()),
259 {
260 let block = Block::new(&[]);
261 let true_block = Block::new(&[(index_type, location)]);
262 let false_block = Block::new(&[(index_type, location)]);
263
264 let operand = block
265 .append_operation(index::constant(
266 &context,
267 IntegerAttribute::new(index_type, 1),
268 location,
269 ))
270 .result(0)
271 .unwrap()
272 .into();
273
274 let condition = block
275 .append_operation(index::cmp(
276 &context,
277 CmpiPredicate::Eq,
278 operand,
279 operand,
280 location,
281 ))
282 .result(0)
283 .unwrap()
284 .into();
285
286 block.append_operation(cond_br(
287 &context,
288 condition,
289 &true_block,
290 &false_block,
291 &[operand],
292 &[operand],
293 location,
294 ));
295
296 true_block.append_operation(func::r#return(&[], location));
297 false_block.append_operation(func::r#return(&[], location));
298
299 let region = Region::new();
300 region.append_block(block);
301 region.append_block(true_block);
302 region.append_block(false_block);
303 region
304 },
305 &[],
306 location,
307 ));
308
309 assert!(module.as_operation().verify());
310 insta::assert_snapshot!(module.as_operation());
311 }
312
313 #[test]
314 fn compile_switch() {
315 let context = Context::new();
316 load_all_dialects(&context);
317
318 let location = Location::unknown(&context);
319 let module = Module::new(location);
320 let i32_type: Type = IntegerType::new(&context, 32).into();
321
322 module.body().append_operation(func::func(
323 &context,
324 StringAttribute::new(&context, "foo"),
325 TypeAttribute::new(FunctionType::new(&context, &[], &[]).into()),
326 {
327 let block = Block::new(&[]);
328 let default_block = Block::new(&[(i32_type, location)]);
329 let first_block = Block::new(&[(i32_type, location)]);
330 let second_block = Block::new(&[(i32_type, location)]);
331
332 let operand = block
333 .append_operation(arith::constant(
334 &context,
335 IntegerAttribute::new(i32_type, 1).into(),
336 location,
337 ))
338 .result(0)
339 .unwrap()
340 .into();
341
342 block.append_operation(
343 switch(
344 &context,
345 &[0, 1],
346 operand,
347 i32_type,
348 (&default_block, &[operand]),
349 &[(&first_block, &[operand]), (&second_block, &[operand])],
350 location,
351 )
352 .unwrap(),
353 );
354
355 default_block.append_operation(func::r#return(&[], location));
356 first_block.append_operation(func::r#return(&[], location));
357 second_block.append_operation(func::r#return(&[], location));
358
359 let region = Region::new();
360 region.append_block(block);
361 region.append_block(default_block);
362 region.append_block(first_block);
363 region.append_block(second_block);
364 region
365 },
366 &[],
367 location,
368 ));
369
370 assert!(module.as_operation().verify());
371 insta::assert_snapshot!(module.as_operation());
372 }
373}