1use cubecl_macros_internal::NamedRewrite;
9use itertools::Itertools;
10use pliron::{
11 attribute::AttrObj,
12 basic_block::BasicBlock,
13 builtin::attributes::IntegerAttr,
14 combine::{
15 Parser, optional,
16 parser::char::{self, spaces},
17 },
18 identifier::Identifier,
19 input_err,
20 irbuild::inserter::OpInsertionPoint,
21 irfmt::parsers::{delimited_list_parser, process_parsed_ssa_defs, spaced, ssa_opd_parser},
22 linked_list::ContainsLinkedList,
23 location::Location,
24 op::{OpBox, OpObj},
25 opts::mem2reg::AllocInfo,
26 parsable::{IntoParseResult, Parsable, ParseResult, StateStream},
27 printable::Printable,
28 region::Region,
29 utils::table::{HMap, SmallMap},
30 value::Use,
31};
32
33use crate::{
34 RecursiveMemoryEffects, RecursiveSideEffects, RecursivelySpeculatable,
35 attributes::{BoolAttr, IntegerVecAttr, ZeroAttr},
36 dialect::{
37 BlockPtrExt,
38 branch::{self, ConditionOp, YieldOp},
39 },
40 interfaces::{
41 CanonicalizeInterface,
42 control_flow::{
43 InvocationBounds, RegionBranchOpInterface, RegionPredecessor, RegionSuccessor,
44 },
45 memory_slot::{
46 MemoryRegionPredecessor, MemorySSAContext, MemorySSARegionOpInterface, MemoryValue,
47 PromotableRegionOpInterface, RegionMemoryPhiInputs, RegionMemoryValue,
48 },
49 uniformity::{UniformRegionOpInterface, Uniformity},
50 },
51 prelude::*,
52 small_map,
53 types::scalar::BoolType,
54};
55
56#[pliron_op(
57 name = "scf.if",
58 format = "$0 ` : ` types(CharSpace(`,`)) ` then ` region($0) ` else ` region($1)",
59 verifier = "succ"
60)]
61#[op_interfaces(NOpdsInterface<1>, NRegionsInterface<2>, SingleBlockRegionInterface, OperandNOfType<0, BoolType>)]
62#[op_traits(RecursiveMemoryEffects, RecursiveSideEffects, RecursivelySpeculatable)]
63pub struct IfOp;
64
65impl IfOp {
66 pub fn new(ctx: &mut Context, results: Vec<TypeHandle>, cond: Value) -> Self {
67 let op = Operation::new(
68 ctx,
69 Self::get_concrete_op_info(),
70 results,
71 vec![cond],
72 vec![],
73 2,
74 );
75
76 let then_region = op.deref_mut(ctx).get_region(0);
77 let then_body = BasicBlock::new(ctx, Some("then".try_into().unwrap()), vec![]);
78 then_body.insert_at_front(then_region, ctx);
79
80 let else_region = op.deref_mut(ctx).get_region(1);
81 let else_body = BasicBlock::new(ctx, Some("else".try_into().unwrap()), vec![]);
82 else_body.insert_at_front(else_region, ctx);
83
84 Self { op }
85 }
86
87 pub fn condition(&self, ctx: &Context) -> Value {
88 self.get_operation().deref(ctx).get_operand(0)
89 }
90
91 pub fn then_region(&self, ctx: &Context) -> Ptr<Region> {
92 self.get_operation().deref(ctx).get_region(0)
93 }
94
95 pub fn then_block(&self, ctx: &Context) -> Ptr<BasicBlock> {
96 self.get_body(ctx, 0)
97 }
98
99 pub fn else_region(&self, ctx: &Context) -> Ptr<Region> {
100 self.get_operation().deref(ctx).get_region(1)
101 }
102
103 pub fn else_block(&self, ctx: &Context) -> Ptr<BasicBlock> {
104 self.get_body(ctx, 1)
105 }
106
107 pub fn get_result(&self, ctx: &Context, res_idx: usize) -> Value {
108 self.get_operation().deref(ctx).get_result(res_idx)
109 }
110
111 pub fn results(&self, ctx: &Context) -> Vec<Value> {
112 self.get_operation().results(ctx)
113 }
114
115 pub fn result_types(&self, ctx: &Context) -> Vec<TypeHandle> {
116 self.get_operation().result_types(ctx)
117 }
118}
119
120fn inline_block(
121 ctx: &Context,
122 rewriter: &mut dyn Rewriter,
123 block: Ptr<BasicBlock>,
124 insertion_point: OpInsertionPoint,
125) {
126 let ops = block.deref(ctx).iter(ctx).collect::<Vec<_>>();
127 let mut insertion_pt = insertion_point;
128 for op in ops {
129 if !op.is_terminator(ctx) {
130 rewriter.move_operation(ctx, op, insertion_pt);
131 insertion_pt = OpInsertionPoint::AfterOperation(op);
132 }
133 }
134}
135
136#[op_interface_impl]
137impl BranchToSCFOp for branch::IfOp {
138 fn to_scf(
139 &self,
140 ctx: &mut Context,
141 rewriter: &mut DialectConversionRewriter,
142 _: &OperandsInfo,
143 ) -> Result<()> {
144 let opds = self.get_operation().operands(ctx);
145 let op = Operation::new(ctx, IfOp::get_concrete_op_info(), vec![], opds, vec![], 0);
146
147 let regions = self.get_operation().regions(ctx);
148 for region in regions {
149 Region::move_to_op(region, op, ctx);
150 }
151
152 rewriter.append_operation(ctx, op);
153 rewriter.replace_operation(ctx, self.get_operation(), op);
154 Ok(())
155 }
156}
157
158#[op_interface_impl]
159impl PromotableRegionOpInterface for IfOp {
160 fn is_region_promotable(&self, _: &Context, _: &AllocInfo, _: Ptr<Region>, _: bool) -> bool {
161 true
162 }
163
164 fn setup_promotion(
165 &self,
166 ctx: &mut Context,
167 _: &AllocInfo,
168 reaching_def: Value,
169 _: bool,
170 regions_to_process: &mut SmallMap<Ptr<Region>, Value, 2>,
171 ) {
172 regions_to_process.insert(self.then_region(ctx), reaching_def);
173 regions_to_process.insert(self.else_region(ctx), reaching_def);
174 }
175
176 fn finalize_promotion(
177 &self,
178 ctx: &mut Context,
179 alloc: &AllocInfo,
180 entry_reaching_def: Value,
181 has_value_stores: bool,
182 reaching_at_block_end: &HMap<Ptr<BasicBlock>, Value>,
183 ) -> Value {
184 if !has_value_stores {
185 return entry_reaching_def;
186 }
187
188 update_terminator::<YieldOp>(
189 ctx,
190 self.then_block(ctx),
191 entry_reaching_def,
192 reaching_at_block_end,
193 );
194 update_terminator::<YieldOp>(
195 ctx,
196 self.else_block(ctx),
197 entry_reaching_def,
198 reaching_at_block_end,
199 );
200
201 let res_idx = Operation::push_result(self.get_operation(), ctx, alloc.ty);
202 self.get_result(ctx, res_idx)
203 }
204}
205
206#[op_interface_impl]
207impl MemorySSARegionOpInterface for IfOp {
208 fn setup_memory_ssa(
209 &self,
210 ctx: &Context,
211 _state: &mut MemorySSAContext,
212 reaching_def: MemoryValue,
213 _has_memory_defs: bool,
214 regions_to_process: &mut SmallMap<Ptr<Region>, MemoryValue, 2>,
215 ) {
216 regions_to_process.insert(self.then_region(ctx), reaching_def);
217 regions_to_process.insert(self.else_region(ctx), reaching_def);
218 }
219
220 fn finalize_memory_ssa(
221 &self,
222 ctx: &Context,
223 _state: &mut MemorySSAContext,
224 entry_reaching_def: MemoryValue,
225 has_memory_defs: bool,
226 _reaching_at_region_entry: &HMap<Ptr<Region>, MemoryValue>,
227 reaching_at_block_end: &HMap<Ptr<BasicBlock>, MemoryValue>,
228 _region_phis: &mut SmallMap<Ptr<Region>, RegionMemoryPhiInputs, 2>,
229 ) -> RegionMemoryValue {
230 if !has_memory_defs {
231 return RegionMemoryValue::Forward(entry_reaching_def);
232 }
233
234 let (then_pred, reaching_then) = block_mem_val(
235 self.then_block(ctx),
236 entry_reaching_def,
237 reaching_at_block_end,
238 );
239 let (else_pred, reaching_else) = block_mem_val(
240 self.else_block(ctx),
241 entry_reaching_def,
242 reaching_at_block_end,
243 );
244
245 RegionMemoryValue::RegionPhi(small_map! {
246 then_pred => reaching_then,
247 else_pred => reaching_else
248 })
249 }
250}
251
252#[op_interface_impl]
253impl RegionBranchOpInterface for IfOp {
254 fn entry_successor_regions(
255 &self,
256 ctx: &Context,
257 operands: &[Option<AttrObj>],
258 ) -> Vec<RegionSuccessor> {
259 let Some(attr) = operands[0].as_ref() else {
260 return self.successor_regions(ctx, RegionPredecessor::Parent);
261 };
262 let zero = attr.downcast_ref::<ZeroAttr>().map(|_| false);
263 let bool = attr.downcast_ref::<BoolAttr>().map(|it| it.0);
264 let Some(const_cond) = zero.or(bool) else {
265 return self.successor_regions(ctx, RegionPredecessor::Parent);
266 };
267 match const_cond {
268 true => vec![self.then_region(ctx).into()],
269 false => vec![self.else_region(ctx).into()],
270 }
271 }
272
273 fn successor_regions(&self, ctx: &Context, pred: RegionPredecessor) -> Vec<RegionSuccessor> {
274 match pred {
275 RegionPredecessor::Parent => {
276 vec![self.then_region(ctx).into(), self.else_region(ctx).into()]
277 }
278 RegionPredecessor::Terminator(_) => {
279 vec![RegionSuccessor::AfterOp]
280 }
281 }
282 }
283
284 fn successor_inputs(&self, ctx: &Context, successor: RegionSuccessor) -> Vec<Value> {
285 match successor {
286 RegionSuccessor::Region(_) => vec![],
287 RegionSuccessor::AfterOp => self.results(ctx),
288 }
289 }
290
291 fn region_invocation_bounds(
292 &self,
293 _ctx: &Context,
294 operands: &[Option<AttrObj>],
295 ) -> Vec<InvocationBounds> {
296 if let Some(cond) = operands[0]
297 .as_ref()
298 .and_then(|it| it.downcast_ref::<BoolAttr>())
299 {
300 match cond.0 {
301 true => vec![InvocationBounds::once(), InvocationBounds::never()],
302 false => vec![InvocationBounds::never(), InvocationBounds::once()],
303 }
304 } else {
305 vec![InvocationBounds::zero_or_one(); 2]
306 }
307 }
308}
309
310#[op_interface_impl]
311impl UniformRegionOpInterface for IfOp {
312 fn result_uniformity(&self, _ctx: &Context, operands: &[Uniformity]) -> Uniformity {
313 operands[0]
314 }
315
316 fn entry_successor_region_uniformity(
317 &self,
318 _ctx: &Context,
319 operands: &[Uniformity],
320 ) -> Vec<Uniformity> {
321 vec![operands[0], operands[0]]
322 }
323}
324
325impl IfOp {
326 fn remove_unused_carried(&self, ctx: &mut Context) -> Result<()> {
327 let results = self.get_operation().results(ctx);
328 for (idx, res) in results.into_iter().enumerate().rev() {
329 if !res.is_used(ctx) {
330 remove_from_terminator::<YieldOp>(ctx, self.then_block(ctx), idx);
331 remove_from_terminator::<YieldOp>(ctx, self.else_block(ctx), idx);
332 Operation::remove_result(self.get_operation(), ctx, idx);
333 }
334 }
335 Ok(())
336 }
337
338 fn fold(&self, ctx: &mut Context, rewriter: &mut MatchRewriter) -> Result<()> {
339 let op = self.get_operation();
340 let operands = const_operands(ctx, op);
341 let valid_branches = self.entry_successor_regions(ctx, &operands);
342 let &[RegionSuccessor::Region(taken)] = valid_branches.as_slice() else {
343 return Ok(());
344 };
345 let taken = taken.deref(ctx).get_entry_block().unwrap();
346
347 let term = taken.deref(ctx).get_terminator(ctx);
348 if let Some(term) = term
349 && term.is_op::<YieldOp>(ctx)
350 {
351 let results = self.results(ctx);
352 let yielded = term.operands(ctx);
353 assert_eq!(results.len(), yielded.len(), "Yield doesn't match results");
354 for (res, yielded) in results.into_iter().zip(yielded) {
355 rewriter.replace_value_uses_with(ctx, res, yielded);
356 Operation::pop_operand(term, ctx);
357 }
358 }
359
360 inline_block(ctx, rewriter, taken, OpInsertionPoint::BeforeOperation(op));
361 rewriter.erase_operation(ctx, op);
362
363 Ok(())
364 }
365}
366
367#[op_interface_impl]
368impl CanonicalizeInterface for IfOp {
369 fn canonicalize(&self, ctx: &mut Context, rewriter: &mut MatchRewriter) -> Result<()> {
370 self.remove_unused_carried(ctx)?;
371 self.fold(ctx, rewriter)?;
372 Ok(())
373 }
374}
375
376#[pliron_op(
377 name = "scf.switch",
378 format,
379 attributes = (scf_switch_cases: IntegerVecAttr),
380 verifier = "succ"
381)]
382#[op_interfaces(NOpdsInterface<1>, SingleBlockRegionInterface)]
383#[op_traits(RecursiveMemoryEffects, RecursiveSideEffects, RecursivelySpeculatable)]
384pub struct SwitchOp;
385
386impl SwitchOp {
387 pub fn new(ctx: &mut Context, results: Vec<TypeHandle>, value: Value) -> Self {
388 let op = Operation::new(
389 ctx,
390 Self::get_concrete_op_info(),
391 results,
392 vec![value],
393 vec![],
394 1,
395 );
396
397 let default_region = op.deref_mut(ctx).get_region(0);
398 let default_body = BasicBlock::new(ctx, Some("default".try_into().unwrap()), vec![]);
399 default_body.insert_at_front(default_region, ctx);
400
401 Self { op }
402 }
403
404 pub fn value(&self, ctx: &Context) -> Value {
405 self.get_operation().deref(ctx).get_operand(0)
406 }
407
408 pub fn default_region(&self, ctx: &Context) -> Ptr<Region> {
409 self.get_operation().deref(ctx).get_region(0)
410 }
411
412 pub fn default_block(&self, ctx: &Context) -> Ptr<BasicBlock> {
413 self.get_body(ctx, 0)
414 }
415
416 pub fn case_regions(&self, ctx: &Context) -> Vec<Ptr<Region>> {
417 self.get_operation().deref(ctx).regions().skip(1).collect()
418 }
419
420 pub fn case_blocks(&self, ctx: &Context) -> Vec<Ptr<BasicBlock>> {
421 self.case_regions(ctx)
422 .iter()
423 .map(|reg| reg.deref(ctx).get_entry_block().unwrap())
424 .collect()
425 }
426
427 pub fn append_case_block(&self, ctx: &mut Context) -> Ptr<BasicBlock> {
428 let region = Operation::add_region(self.get_operation(), ctx);
429 let body = BasicBlock::new(ctx, None, vec![]);
430 body.insert_at_front(region, ctx);
431 region.deref(ctx).get_head().unwrap()
432 }
433
434 pub fn cases(&self, ctx: &Context) -> Vec<(IntegerAttr, Ptr<BasicBlock>)> {
435 let cases = self.get_attr_scf_switch_cases(ctx).unwrap().clone().0;
436 let out = (0..cases.len()).map(|i| {
437 let value = cases[i].clone();
438 let block = self.get_body(ctx, i + 1);
439 (value, block)
440 });
441 out.collect()
442 }
443
444 pub fn cases_regions(&self, ctx: &Context) -> Vec<(IntegerAttr, Ptr<Region>)> {
445 let cases = self.get_attr_scf_switch_cases(ctx).unwrap().clone().0;
446 let out = (0..cases.len()).map(|i| {
447 let value = cases[i].clone();
448 let block = self.get_operation().deref(ctx).get_region(i + 1);
449 (value, block)
450 });
451 out.collect()
452 }
453
454 pub fn cases_values(&self, ctx: &Context) -> Vec<IntegerAttr> {
455 self.get_attr_scf_switch_cases(ctx).unwrap().0.clone()
456 }
457
458 pub fn get_case_destinations(&self, ctx: &Context) -> Vec<Ptr<BasicBlock>> {
459 let op = self.get_operation().deref(ctx);
460 (1..op.regions().count())
461 .map(|i| self.get_body(ctx, i))
462 .collect()
463 }
464
465 pub fn set_attr_cases(&self, ctx: &Context, cases: impl IntoIterator<Item = IntegerAttr>) {
466 self.set_attr_scf_switch_cases(ctx, IntegerVecAttr(cases.into_iter().collect()));
467 }
468
469 pub fn results(&self, ctx: &Context) -> Vec<Value> {
470 self.get_operation().results(ctx)
471 }
472
473 pub fn result_types(&self, ctx: &Context) -> Vec<TypeHandle> {
474 self.get_operation().result_types(ctx)
475 }
476}
477
478#[op_interface_impl]
479impl BranchToSCFOp for branch::SwitchOp {
480 fn to_scf(
481 &self,
482 ctx: &mut Context,
483 rewriter: &mut DialectConversionRewriter,
484 _: &OperandsInfo,
485 ) -> Result<()> {
486 let opds = self.get_operation().operands(ctx);
487 let info = SwitchOp::get_concrete_op_info();
488 let op = Operation::new(ctx, info, vec![], opds, vec![], 0);
489 let cases = self.get_attr_branch_switch_cases(ctx).unwrap().clone();
490 SwitchOp { op }.set_attr_scf_switch_cases(ctx, cases);
491
492 let regions = self.get_operation().regions(ctx);
493 for region in regions {
494 Region::move_to_op(region, op, ctx);
495 }
496
497 rewriter.append_operation(ctx, op);
498 rewriter.replace_operation(ctx, self.get_operation(), op);
499 Ok(())
500 }
501}
502
503#[op_interface_impl]
504impl PromotableRegionOpInterface for SwitchOp {
505 fn is_region_promotable(&self, _: &Context, _: &AllocInfo, _: Ptr<Region>, _: bool) -> bool {
506 true
507 }
508
509 fn setup_promotion(
510 &self,
511 ctx: &mut Context,
512 _: &AllocInfo,
513 reaching_def: Value,
514 _: bool,
515 regions_to_process: &mut SmallMap<Ptr<Region>, Value, 2>,
516 ) {
517 regions_to_process.insert(self.default_region(ctx), reaching_def);
518 for case_region in self.case_regions(ctx) {
519 regions_to_process.insert(case_region, reaching_def);
520 }
521 }
522
523 fn finalize_promotion(
524 &self,
525 ctx: &mut Context,
526 alloc: &AllocInfo,
527 reaching_def: Value,
528 has_value_stores: bool,
529 reaching_at_block_end: &HMap<Ptr<BasicBlock>, Value>,
530 ) -> Value {
531 if !has_value_stores {
532 return reaching_def;
533 }
534
535 update_terminator::<YieldOp>(
536 ctx,
537 self.default_block(ctx),
538 reaching_def,
539 reaching_at_block_end,
540 );
541 for case_block in self.case_blocks(ctx) {
542 update_terminator::<YieldOp>(ctx, case_block, reaching_def, reaching_at_block_end);
543 }
544
545 let res_idx = Operation::push_result(self.get_operation(), ctx, alloc.ty);
546 self.get_operation().deref(ctx).get_result(res_idx)
547 }
548}
549
550#[op_interface_impl]
551impl MemorySSARegionOpInterface for SwitchOp {
552 fn setup_memory_ssa(
553 &self,
554 ctx: &Context,
555 _state: &mut MemorySSAContext,
556 reaching_def: MemoryValue,
557 _has_memory_defs: bool,
558 regions_to_process: &mut SmallMap<Ptr<Region>, MemoryValue, 2>,
559 ) {
560 regions_to_process.insert(self.default_region(ctx), reaching_def);
561 for case_region in self.case_regions(ctx) {
562 regions_to_process.insert(case_region, reaching_def);
563 }
564 }
565
566 fn finalize_memory_ssa(
567 &self,
568 ctx: &Context,
569 _state: &mut MemorySSAContext,
570 entry_reaching_def: MemoryValue,
571 has_memory_defs: bool,
572 _reaching_at_region_entry: &HMap<Ptr<Region>, MemoryValue>,
573 reaching_at_block_end: &HMap<Ptr<BasicBlock>, MemoryValue>,
574 _region_phis: &mut SmallMap<Ptr<Region>, RegionMemoryPhiInputs, 2>,
575 ) -> RegionMemoryValue {
576 if !has_memory_defs {
577 return RegionMemoryValue::Forward(entry_reaching_def);
578 }
579
580 let mut phi_inputs = SmallMap::new();
581
582 let (default_pred, reaching_default) = block_mem_val(
583 self.default_block(ctx),
584 entry_reaching_def,
585 reaching_at_block_end,
586 );
587 phi_inputs.insert(default_pred, reaching_default);
588
589 for case_block in self.case_blocks(ctx) {
590 let (case_pred, reaching_case) =
591 block_mem_val(case_block, entry_reaching_def, reaching_at_block_end);
592 phi_inputs.insert(case_pred, reaching_case);
593 }
594
595 RegionMemoryValue::RegionPhi(phi_inputs)
596 }
597}
598
599#[op_interface_impl]
600impl RegionBranchOpInterface for SwitchOp {
601 fn entry_successor_regions(
602 &self,
603 ctx: &Context,
604 operands: &[Option<AttrObj>],
605 ) -> Vec<RegionSuccessor> {
606 let Some(attr) = &operands[0] else {
607 return self.successor_regions(ctx, RegionPredecessor::Parent);
608 };
609 let Some(attr) = attr.downcast_ref::<IntegerAttr>() else {
610 return self.successor_regions(ctx, RegionPredecessor::Parent);
611 };
612 if let Some(&(_, case)) = self.cases_regions(ctx).iter().find(|(val, _)| val == attr) {
613 vec![case.into()]
614 } else {
615 vec![self.default_region(ctx).into()]
616 }
617 }
618
619 fn successor_regions(&self, ctx: &Context, pred: RegionPredecessor) -> Vec<RegionSuccessor> {
620 match pred {
621 RegionPredecessor::Parent => {
622 let op = self.get_operation().deref(ctx);
623 op.regions().map(Into::into).collect()
624 }
625 RegionPredecessor::Terminator(_) => {
626 vec![RegionSuccessor::AfterOp]
627 }
628 }
629 }
630
631 fn successor_inputs(&self, ctx: &Context, successor: RegionSuccessor) -> Vec<Value> {
632 match successor {
633 RegionSuccessor::Region(_) => vec![],
634 RegionSuccessor::AfterOp => self.results(ctx),
635 }
636 }
637
638 fn region_invocation_bounds(
639 &self,
640 ctx: &Context,
641 operands: &[Option<AttrObj>],
642 ) -> Vec<InvocationBounds> {
643 let num_regions = self.get_operation().deref(ctx).num_regions();
644 let Some(attr) = operands[0].as_ref() else {
645 return vec![InvocationBounds::zero_or_one(); num_regions];
646 };
647 let Some(attr) = attr.downcast_ref::<IntegerAttr>() else {
648 return vec![InvocationBounds::zero_or_one(); num_regions];
649 };
650
651 let case_idx = self.cases_values(ctx).iter().position(|it| it == attr);
652 let executed_idx = case_idx.map(|i| i + 1).unwrap_or(0);
653 let mut bounds = vec![InvocationBounds::never(); num_regions];
654 bounds[executed_idx] = InvocationBounds::once();
655 bounds
656 }
657}
658
659#[op_interface_impl]
660impl UniformRegionOpInterface for SwitchOp {
661 fn result_uniformity(&self, _ctx: &Context, operands: &[Uniformity]) -> Uniformity {
662 operands[0]
663 }
664
665 fn entry_successor_region_uniformity(
666 &self,
667 ctx: &Context,
668 operands: &[Uniformity],
669 ) -> Vec<Uniformity> {
670 let num_regions = self.get_operation().deref(ctx).num_regions();
671 vec![operands[0]; num_regions]
672 }
673}
674
675impl SwitchOp {
676 fn remove_unused_carried(&self, ctx: &mut Context) -> Result<()> {
677 let results = self.get_operation().results(ctx);
678 for (idx, res) in results.into_iter().enumerate().rev() {
679 if !res.is_used(ctx) {
680 remove_from_terminator::<YieldOp>(ctx, self.default_block(ctx), idx);
681 for case_block in self.case_blocks(ctx) {
682 remove_from_terminator::<YieldOp>(ctx, case_block, idx);
683 }
684 Operation::remove_result(self.get_operation(), ctx, idx);
685 }
686 }
687 Ok(())
688 }
689
690 fn fold(&self, ctx: &mut Context, rewriter: &mut MatchRewriter) -> Result<()> {
691 let op = self.get_operation();
692 let operands = const_operands(ctx, op);
693 let valid_branches = self.entry_successor_regions(ctx, &operands);
694 let &[RegionSuccessor::Region(taken)] = valid_branches.as_slice() else {
695 return Ok(());
696 };
697 let taken = taken.deref(ctx).get_entry_block().unwrap();
698
699 let term = taken.deref(ctx).get_terminator(ctx);
700 if let Some(term) = term
701 && term.is_op::<YieldOp>(ctx)
702 {
703 let results = self.results(ctx);
704 let yielded = term.operands(ctx);
705 assert_eq!(results.len(), yielded.len(), "Yield doesn't match results");
706 for (res, yielded) in results.into_iter().zip(yielded) {
707 rewriter.replace_value_uses_with(ctx, res, yielded);
708 Operation::pop_operand(term, ctx);
709 }
710 }
711
712 inline_block(ctx, rewriter, taken, OpInsertionPoint::BeforeOperation(op));
713 rewriter.erase_operation(ctx, op);
714
715 Ok(())
716 }
717}
718
719#[op_interface_impl]
720impl CanonicalizeInterface for SwitchOp {
721 fn canonicalize(&self, ctx: &mut Context, rewriter: &mut MatchRewriter) -> Result<()> {
722 self.remove_unused_carried(ctx)?;
723 self.fold(ctx, rewriter)?;
724 Ok(())
725 }
726}
727
728#[pliron_op(name = "scf.for", verifier = "succ")]
729#[op_interfaces(OneRegionInterface, SingleBlockRegionInterface)]
730#[op_traits(RecursiveMemoryEffects, RecursiveSideEffects)]
731pub struct ForOp;
732
733impl ForOp {
734 pub fn new(
735 ctx: &mut Context,
736 results: Vec<TypeHandle>,
737 start: Value,
738 end: Value,
739 step: Value,
740 carried_values_init: Vec<Value>,
741 ) -> Self {
742 let iter_ty = start.get_type(ctx);
743 let mut body_args = vec![iter_ty];
744 body_args.extend(carried_values_init.iter().map(|it| it.get_type(ctx)));
745
746 let (operands, segments) =
747 Self::compute_segment_sizes(vec![vec![start, end, step], carried_values_init]);
748
749 let op = Self {
750 op: Operation::new(
751 ctx,
752 Self::get_concrete_op_info(),
753 results,
754 operands,
755 vec![],
756 1,
757 ),
758 };
759
760 op.set_operand_segment_sizes(ctx, segments);
761
762 let body_region = op.loop_region(ctx);
763 let body = BasicBlock::new(ctx, Some("body".try_into().unwrap()), body_args);
764 body.insert_at_front(body_region, ctx);
765
766 op
767 }
768
769 pub fn iter_var(&self, ctx: &Context) -> Value {
770 self.loop_body(ctx).deref(ctx).get_argument(0)
771 }
772
773 pub fn start(&self, ctx: &Context) -> Value {
774 self.get_operation().deref(ctx).get_operand(0)
775 }
776
777 pub fn end(&self, ctx: &Context) -> Value {
778 self.get_operation().deref(ctx).get_operand(1)
779 }
780
781 pub fn step(&self, ctx: &Context) -> Value {
782 self.get_operation().deref(ctx).get_operand(2)
783 }
784
785 pub fn initial_carried_values(&self, ctx: &Context) -> Vec<Value> {
786 self.get_segment(ctx, 1)
787 }
788
789 pub fn push_initial_carried_value(&self, ctx: &mut Context, value: Value) -> usize {
790 self.push_to_segment(ctx, 1, value)
791 }
792
793 pub fn remove_initial_carried_value(&self, ctx: &mut Context, opd_idx: usize) -> Value {
794 self.remove_from_segment(ctx, 1, opd_idx)
795 }
796
797 pub fn carried_values(&self, ctx: &Context) -> Vec<Value> {
798 self.loop_body(ctx).deref(ctx).arguments().skip(1).collect()
799 }
800
801 pub fn get_carried_value(&self, ctx: &Context, arg_idx: usize) -> Value {
802 self.loop_body(ctx).deref(ctx).get_argument(arg_idx + 1)
803 }
804
805 pub fn push_carried_value(&self, ctx: &mut Context, ty: TypeHandle) -> usize {
806 BasicBlock::push_argument(self.loop_body(ctx), ctx, ty)
807 }
808
809 pub fn remove_carried_value(&self, ctx: &Context, arg_idx: usize) {
810 BasicBlock::remove_argument(self.loop_body(ctx), ctx, arg_idx + 1)
811 }
812
813 pub fn loop_region(&self, ctx: &Context) -> Ptr<Region> {
814 self.get_operation().deref(ctx).get_region(0)
815 }
816
817 pub fn loop_body(&self, ctx: &Context) -> Ptr<BasicBlock> {
818 self.get_body(ctx, 0)
819 }
820
821 pub fn results(&self, ctx: &Context) -> Vec<Value> {
822 self.get_operation().results(ctx)
823 }
824
825 pub fn result_types(&self, ctx: &Context) -> Vec<TypeHandle> {
826 self.get_operation().result_types(ctx)
827 }
828}
829
830#[op_interface_impl]
831impl OperandSegmentInterface for ForOp {
832 fn expected_num_segments(&self, _: &Context) -> Option<usize> {
833 Some(2)
834 }
835}
836
837impl Printable for ForOp {
838 fn fmt(
839 &self,
840 ctx: &Context,
841 state: &pliron::printable::State,
842 f: &mut core::fmt::Formatter<'_>,
843 ) -> core::fmt::Result {
844 let results = self.results(ctx);
845 if !results.is_empty() {
846 let results = results.iter().map(|it| it.disp(ctx)).join(", ");
847 write!(f, "{results} = ")?;
848 }
849 let args = self.initial_carried_values(ctx);
850 write!(
851 f,
852 "{} {} to {} step {} ",
853 self.get_opid(),
854 self.start(ctx).disp(ctx),
855 self.end(ctx).disp(ctx),
856 self.step(ctx).disp(ctx),
857 )?;
858
859 if !args.is_empty() {
860 let args = args.iter().map(|it| it.disp(ctx)).join(", ");
861 write!(f, "iter_args({args}) ")?;
862 }
863
864 self.loop_region(ctx).fmt(ctx, state, f)?;
865
866 Ok(())
867 }
868}
869
870impl Parsable for ForOp {
871 type Arg = Vec<(Identifier, Location)>;
872 type Parsed = OpObj;
873
874 fn parse<'a>(
875 state_stream: &mut StateStream<'a>,
876 results: Self::Arg,
877 ) -> ParseResult<'a, Self::Parsed> {
878 let cur_loc = state_stream.loc();
879 let iter_args = delimited_list_parser('(', ')', ',', ssa_opd_parser()).skip(spaces());
880 let mut iter_args_parser =
881 optional(spaced(char::string("iter_args").with(iter_args))).skip(spaces());
882 let mut range_parser = (
883 ssa_opd_parser().skip(spaced(char::string("to"))),
884 ssa_opd_parser().skip(spaced(char::string("step"))),
885 ssa_opd_parser(),
886 );
887
888 let (start, end, step) = range_parser.parse_stream(state_stream).into_result()?.0;
889 let args = iter_args_parser.parse_stream(state_stream).into_result()?.0;
890 let args = args.unwrap_or_default();
891
892 if args.len() != results.len() {
893 input_err!(
894 cur_loc,
895 "expected {} results based on iter vars, got {} during parsing",
896 args.len(),
897 results.len()
898 )?;
899 }
900
901 let result_types = args
902 .iter()
903 .map(|it| it.get_type(state_stream.state.ctx))
904 .collect();
905 let (operands, segments) = Self::compute_segment_sizes(vec![vec![start, end, step], args]);
906
907 let op = Operation::new(
908 state_stream.state.ctx,
909 Self::get_concrete_op_info(),
910 result_types,
911 operands,
912 vec![],
913 0,
914 );
915 process_parsed_ssa_defs(state_stream, &results, op)?;
916
917 Region::parse(state_stream, op)?;
918 let op = ForOp { op };
919 op.set_operand_segment_sizes(state_stream.state.ctx, segments);
920
921 Ok(OpBox::new(op)).into_parse_result()
922 }
923}
924
925#[op_interface_impl]
926impl BranchToSCFOp for branch::RangeLoopOp {
927 fn to_scf(
928 &self,
929 ctx: &mut Context,
930 rewriter: &mut DialectConversionRewriter,
931 _: &OperandsInfo,
932 ) -> Result<()> {
933 let opds = self.get_operation().operands(ctx);
934 let (operands, segments) = ForOp::compute_segment_sizes(vec![opds, vec![]]);
935 let info = ForOp::get_concrete_op_info();
936 let op = Operation::new(ctx, info, vec![], operands, vec![], 0);
937
938 let regions = self.get_operation().regions(ctx);
939 for region in regions {
940 Region::move_to_op(region, op, ctx);
941 }
942
943 rewriter.append_operation(ctx, op);
944 rewriter.replace_operation(ctx, self.get_operation(), op);
945
946 let op = ForOp { op };
947 op.set_operand_segment_sizes(ctx, segments);
948
949 Ok(())
950 }
951}
952
953#[op_interface_impl]
954impl PromotableRegionOpInterface for ForOp {
955 fn is_region_promotable(&self, _: &Context, _: &AllocInfo, _: Ptr<Region>, _: bool) -> bool {
956 true
957 }
958
959 fn setup_promotion(
960 &self,
961 ctx: &mut Context,
962 alloc: &AllocInfo,
963 reaching_def: Value,
964 has_value_stores: bool,
965 regions_to_process: &mut SmallMap<Ptr<Region>, Value, 2>,
966 ) {
967 let body_region = self.loop_region(ctx);
968 if !has_value_stores {
969 regions_to_process.insert(body_region, reaching_def);
970 return;
971 }
972
973 self.push_initial_carried_value(ctx, reaching_def);
974 let idx = BasicBlock::push_argument(self.loop_body(ctx), ctx, alloc.ty);
975 let new_arg = self.loop_body(ctx).deref(ctx).get_argument(idx);
976 regions_to_process.insert(body_region, new_arg);
977 }
978
979 fn finalize_promotion(
980 &self,
981 ctx: &mut Context,
982 alloc: &AllocInfo,
983 entry_reaching_def: Value,
984 has_value_stores: bool,
985 reaching_at_block_end: &HMap<Ptr<BasicBlock>, Value>,
986 ) -> Value {
987 if !has_value_stores {
988 return entry_reaching_def;
989 }
990
991 update_terminator::<YieldOp>(
992 ctx,
993 self.loop_body(ctx),
994 entry_reaching_def,
995 reaching_at_block_end,
996 );
997
998 let idx = Operation::push_result(self.get_operation(), ctx, alloc.ty);
999 self.get_operation().deref(ctx).get_result(idx)
1000 }
1001}
1002
1003#[op_interface_impl]
1004impl MemorySSARegionOpInterface for ForOp {
1005 fn setup_memory_ssa(
1006 &self,
1007 ctx: &Context,
1008 state: &mut MemorySSAContext,
1009 reaching_def: MemoryValue,
1010 has_memory_defs: bool,
1011 regions_to_process: &mut SmallMap<Ptr<Region>, MemoryValue, 2>,
1012 ) {
1013 let body_region = self.loop_region(ctx);
1014 if !has_memory_defs {
1015 regions_to_process.insert(body_region, reaching_def);
1016 return;
1017 }
1018
1019 let new_arg = state.new_value_in_block(self.loop_body(ctx));
1020 regions_to_process.insert(body_region, new_arg);
1021 }
1022
1023 fn finalize_memory_ssa(
1024 &self,
1025 ctx: &Context,
1026 _state: &mut MemorySSAContext,
1027 entry_reaching_def: MemoryValue,
1028 has_memory_defs: bool,
1029 _reaching_at_region_entry: &HMap<Ptr<Region>, MemoryValue>,
1030 reaching_at_block_end: &HMap<Ptr<BasicBlock>, MemoryValue>,
1031 region_phis: &mut SmallMap<Ptr<Region>, RegionMemoryPhiInputs, 2>,
1032 ) -> RegionMemoryValue {
1033 let body_region = self.loop_region(ctx);
1034 if !has_memory_defs {
1035 return RegionMemoryValue::Forward(entry_reaching_def);
1036 }
1037
1038 let (body_pred, reaching_body) = block_mem_val(
1039 self.loop_body(ctx),
1040 entry_reaching_def,
1041 reaching_at_block_end,
1042 );
1043
1044 let phi_inputs = small_map! {
1045 MemoryRegionPredecessor::Parent => entry_reaching_def,
1046 body_pred => reaching_body
1047 };
1048
1049 region_phis.insert(body_region, phi_inputs.clone());
1050 RegionMemoryValue::RegionPhi(phi_inputs)
1051 }
1052}
1053
1054#[op_interface_impl]
1055impl RegionBranchOpInterface for ForOp {
1056 fn entry_successor_operands(&self, ctx: &Context, _: RegionSuccessor) -> Vec<Use<Value>> {
1057 operand_range_to_uses(ctx, self.get_operation(), self.segment_range(ctx, 1))
1058 }
1059
1060 fn successor_regions(&self, ctx: &Context, _pred: RegionPredecessor) -> Vec<RegionSuccessor> {
1061 vec![self.loop_region(ctx).into(), RegionSuccessor::AfterOp]
1063 }
1064
1065 fn successor_inputs(&self, ctx: &Context, successor: RegionSuccessor) -> Vec<Value> {
1066 match successor {
1067 RegionSuccessor::Region(_) => self.carried_values(ctx),
1068 RegionSuccessor::AfterOp => self.results(ctx),
1069 }
1070 }
1071}
1072
1073#[op_interface_impl]
1074impl UniformRegionOpInterface for ForOp {
1075 fn result_uniformity(&self, _ctx: &Context, operands: &[Uniformity]) -> Uniformity {
1076 operands[0].min(operands[1]).min(operands[2])
1077 }
1078
1079 fn entry_successor_region_uniformity(
1080 &self,
1081 _ctx: &Context,
1082 operands: &[Uniformity],
1083 ) -> Vec<Uniformity> {
1084 let uniformity = operands[0].min(operands[1]).min(operands[2]);
1085 vec![uniformity, uniformity]
1086 }
1087}
1088
1089#[op_interface_impl]
1090impl CanonicalizeInterface for ForOp {
1091 fn canonicalize(&self, ctx: &mut Context, _rewriter: &mut MatchRewriter) -> Result<()> {
1092 let results = self.get_operation().results(ctx);
1093 let body_block = self.loop_body(ctx);
1094 for (idx, res) in results.into_iter().enumerate().rev() {
1095 let carried = self.get_carried_value(ctx, idx);
1096 if !res.is_used(ctx) && only_used_for_forward::<YieldOp>(ctx, body_block, idx, carried)
1097 {
1098 remove_from_terminator::<YieldOp>(ctx, body_block, idx);
1099 self.remove_initial_carried_value(ctx, idx);
1100 self.remove_carried_value(ctx, idx);
1101 Operation::remove_result(self.get_operation(), ctx, idx);
1102 }
1103 }
1104 Ok(())
1105 }
1106}
1107
1108#[pliron_op(
1109 name = "scf.while",
1110 format = "operands(CharSpace(`,`)) ` : ` types(CharSpace(`,`)) ` ` region($0) ` do ` region($1)",
1111 verifier = "succ"
1112)]
1113#[op_interfaces(NRegionsInterface<2>, SingleBlockRegionInterface)]
1114#[op_traits(RecursiveMemoryEffects, RecursiveSideEffects)]
1115pub struct WhileOp;
1116
1117impl WhileOp {
1118 pub fn new(
1119 ctx: &mut Context,
1120 results: Vec<TypeHandle>,
1121 carried_values_init: Vec<Value>,
1122 ) -> Self {
1123 let carried_types = carried_values_init
1124 .iter()
1125 .map(|it| it.get_type(ctx))
1126 .collect::<Vec<_>>();
1127
1128 let op = Operation::new(
1129 ctx,
1130 Self::get_concrete_op_info(),
1131 results,
1132 carried_values_init,
1133 vec![],
1134 2,
1135 );
1136
1137 let before_region = op.deref_mut(ctx).get_region(0);
1138 let before_block = BasicBlock::new(
1139 ctx,
1140 Some("before".try_into().unwrap()),
1141 carried_types.clone(),
1142 );
1143 before_block.insert_at_front(before_region, ctx);
1144
1145 let after_region = op.deref_mut(ctx).get_region(1);
1146 let after_block = BasicBlock::new(ctx, Some("after".try_into().unwrap()), carried_types);
1147 after_block.insert_at_front(after_region, ctx);
1148
1149 Self { op }
1150 }
1151
1152 pub fn initial_carried_values(&self, ctx: &Context) -> Vec<Value> {
1153 self.get_operation().operands(ctx)
1154 }
1155
1156 pub fn push_initial_carried_value(&self, ctx: &mut Context, value: Value) -> usize {
1157 Operation::push_operand(self.get_operation(), ctx, value)
1158 }
1159
1160 pub fn remove_initial_carried_value(&self, ctx: &mut Context, opd_idx: usize) -> Value {
1161 Operation::remove_operand(self.get_operation(), ctx, opd_idx)
1162 }
1163
1164 pub fn get_before_carried_value(&self, ctx: &Context, arg_idx: usize) -> Value {
1165 self.before_block(ctx).deref(ctx).get_argument(arg_idx)
1166 }
1167
1168 pub fn push_before_carried_value(&self, ctx: &mut Context, ty: TypeHandle) -> usize {
1169 BasicBlock::push_argument(self.before_block(ctx), ctx, ty)
1170 }
1171
1172 pub fn remove_before_carried_value(&self, ctx: &mut Context, arg_idx: usize) {
1173 BasicBlock::remove_argument(self.before_block(ctx), ctx, arg_idx)
1174 }
1175
1176 pub fn get_after_carried_value(&self, ctx: &Context, arg_idx: usize) -> Value {
1177 self.after_block(ctx).deref(ctx).get_argument(arg_idx)
1178 }
1179
1180 pub fn push_after_carried_value(&self, ctx: &mut Context, ty: TypeHandle) -> usize {
1181 BasicBlock::push_argument(self.after_block(ctx), ctx, ty)
1182 }
1183
1184 pub fn remove_after_carried_value(&self, ctx: &mut Context, arg_idx: usize) {
1185 BasicBlock::remove_argument(self.after_block(ctx), ctx, arg_idx)
1186 }
1187
1188 pub fn before_region(&self, ctx: &Context) -> Ptr<Region> {
1189 self.get_operation().deref(ctx).get_region(0)
1190 }
1191
1192 pub fn before_block(&self, ctx: &Context) -> Ptr<BasicBlock> {
1193 self.get_body(ctx, 0)
1194 }
1195
1196 pub fn after_region(&self, ctx: &Context) -> Ptr<Region> {
1197 self.get_operation().deref(ctx).get_region(1)
1198 }
1199
1200 pub fn after_block(&self, ctx: &Context) -> Ptr<BasicBlock> {
1201 self.get_body(ctx, 1)
1202 }
1203
1204 pub fn results(&self, ctx: &Context) -> Vec<Value> {
1205 self.get_operation().results(ctx)
1206 }
1207
1208 pub fn result_types(&self, ctx: &Context) -> Vec<TypeHandle> {
1209 self.get_operation().result_types(ctx)
1210 }
1211}
1212
1213#[op_interface_impl]
1214impl BranchToSCFOp for branch::WhileOp {
1215 fn to_scf(
1216 &self,
1217 ctx: &mut Context,
1218 rewriter: &mut DialectConversionRewriter,
1219 _: &OperandsInfo,
1220 ) -> Result<()> {
1221 let opds = self.get_operation().operands(ctx);
1222 let info = WhileOp::get_concrete_op_info();
1223 let op = Operation::new(ctx, info, vec![], opds, vec![], 0);
1224
1225 let regions = self.get_operation().regions(ctx);
1226 for region in regions {
1227 Region::move_to_op(region, op, ctx);
1228 }
1229
1230 rewriter.append_operation(ctx, op);
1231 rewriter.replace_operation(ctx, self.get_operation(), op);
1232
1233 Ok(())
1234 }
1235}
1236
1237#[op_interface_impl]
1238impl PromotableRegionOpInterface for WhileOp {
1239 fn is_region_promotable(&self, _: &Context, _: &AllocInfo, _: Ptr<Region>, _: bool) -> bool {
1240 true
1241 }
1242
1243 fn setup_promotion(
1244 &self,
1245 ctx: &mut Context,
1246 alloc: &AllocInfo,
1247 reaching_def: Value,
1248 has_value_stores: bool,
1249 regions_to_process: &mut SmallMap<Ptr<Region>, Value, 2>,
1250 ) {
1251 let before_region = self.before_region(ctx);
1252 let after_region = self.after_region(ctx);
1253 if !has_value_stores {
1254 regions_to_process.insert(before_region, reaching_def);
1255 regions_to_process.insert(after_region, reaching_def);
1256 return;
1257 }
1258
1259 self.push_initial_carried_value(ctx, reaching_def);
1260
1261 let idx = BasicBlock::push_argument(self.before_block(ctx), ctx, alloc.ty);
1262 let new_arg = self.before_block(ctx).deref(ctx).get_argument(idx);
1263 regions_to_process.insert(before_region, new_arg);
1264
1265 let idx = BasicBlock::push_argument(self.after_block(ctx), ctx, alloc.ty);
1266 let new_arg = self.after_block(ctx).deref(ctx).get_argument(idx);
1267 regions_to_process.insert(after_region, new_arg);
1268 }
1269
1270 fn finalize_promotion(
1271 &self,
1272 ctx: &mut Context,
1273 alloc: &AllocInfo,
1274 reaching_def: Value,
1275 has_value_stores: bool,
1276 reaching_at_block_end: &HMap<Ptr<BasicBlock>, Value>,
1277 ) -> Value {
1278 if !has_value_stores {
1279 return reaching_def;
1280 }
1281
1282 let before = self.before_block(ctx);
1283 let last_arg = before.deref(ctx).arguments().last();
1284 update_terminator::<ConditionOp>(ctx, before, last_arg.unwrap(), reaching_at_block_end);
1285
1286 let after = self.after_block(ctx);
1287 let last_arg = after.deref(ctx).arguments().last();
1288 update_terminator::<YieldOp>(ctx, after, last_arg.unwrap(), reaching_at_block_end);
1289
1290 let idx = Operation::push_result(self.get_operation(), ctx, alloc.ty);
1291 self.get_operation().deref(ctx).get_result(idx)
1292 }
1293}
1294
1295#[op_interface_impl]
1296impl MemorySSARegionOpInterface for WhileOp {
1297 fn setup_memory_ssa(
1298 &self,
1299 ctx: &Context,
1300 state: &mut MemorySSAContext,
1301 reaching_def: MemoryValue,
1302 has_memory_defs: bool,
1303 regions_to_process: &mut SmallMap<Ptr<Region>, MemoryValue, 2>,
1304 ) {
1305 let before_region = self.before_region(ctx);
1306 let after_region = self.after_region(ctx);
1307 if !has_memory_defs {
1308 regions_to_process.insert(before_region, reaching_def);
1309 regions_to_process.insert(after_region, reaching_def);
1310 return;
1311 }
1312
1313 let new_arg = state.new_value_in_block(self.before_block(ctx));
1314 regions_to_process.insert(before_region, new_arg);
1315
1316 let new_arg = state.new_value_in_block(self.after_block(ctx));
1317 regions_to_process.insert(after_region, new_arg);
1318 }
1319
1320 fn finalize_memory_ssa(
1321 &self,
1322 ctx: &Context,
1323 _state: &mut MemorySSAContext,
1324 entry_reaching_def: MemoryValue,
1325 has_memory_defs: bool,
1326 reaching_at_region_entry: &HMap<Ptr<Region>, MemoryValue>,
1327 reaching_at_block_end: &HMap<Ptr<BasicBlock>, MemoryValue>,
1328 region_phis: &mut SmallMap<Ptr<Region>, RegionMemoryPhiInputs, 2>,
1329 ) -> RegionMemoryValue {
1330 if !has_memory_defs {
1331 return RegionMemoryValue::Forward(entry_reaching_def);
1332 }
1333
1334 let before_region = self.before_region(ctx);
1335 let after_region = self.after_region(ctx);
1336
1337 let arg = reaching_at_region_entry[&before_region];
1338 let (before_pred, reaching_before) =
1339 block_mem_val(self.before_block(ctx), arg, reaching_at_block_end);
1340
1341 let arg = reaching_at_region_entry[&after_region];
1342 let (after_pred, reaching_after) =
1343 block_mem_val(self.after_block(ctx), arg, reaching_at_block_end);
1344
1345 let inputs_before = small_map! {
1346 MemoryRegionPredecessor::Parent => entry_reaching_def,
1347 after_pred => reaching_after
1348 };
1349 let inputs_after = small_map!(before_pred => reaching_before);
1350
1351 region_phis.insert(before_region, inputs_before);
1352 region_phis.insert(after_region, inputs_after);
1353
1354 RegionMemoryValue::Forward(reaching_before)
1355 }
1356}
1357
1358#[op_interface_impl]
1359impl RegionBranchOpInterface for WhileOp {
1360 fn entry_successor_operands(&self, ctx: &Context, _: RegionSuccessor) -> Vec<Use<Value>> {
1361 self.get_operation().deref(ctx).operands_as_uses().collect()
1362 }
1363
1364 fn successor_regions(&self, ctx: &Context, pred: RegionPredecessor) -> Vec<RegionSuccessor> {
1365 match pred {
1366 RegionPredecessor::Parent => vec![self.before_region(ctx).into()],
1367 RegionPredecessor::Terminator(term) => {
1368 let op = term.get_operation();
1369 let parent = op.deref(ctx).get_parent_region(ctx).unwrap();
1370 if parent == self.after_region(ctx) {
1371 vec![self.before_region(ctx).into()]
1372 } else {
1373 vec![RegionSuccessor::AfterOp, self.after_region(ctx).into()]
1374 }
1375 }
1376 }
1377 }
1378
1379 fn successor_inputs(&self, ctx: &Context, successor: RegionSuccessor) -> Vec<Value> {
1380 match successor {
1381 RegionSuccessor::Region(region) if region == self.before_region(ctx) => {
1382 self.before_block(ctx).arguments(ctx)
1383 }
1384 RegionSuccessor::Region(_) => self.after_block(ctx).arguments(ctx),
1385 RegionSuccessor::AfterOp => self.results(ctx),
1386 }
1387 }
1388}
1389
1390#[op_interface_impl]
1391impl UniformRegionOpInterface for WhileOp {
1392 fn result_uniformity(&self, _ctx: &Context, _operands: &[Uniformity]) -> Uniformity {
1393 Uniformity::Cube
1394 }
1395
1396 fn entry_successor_region_uniformity(
1397 &self,
1398 _ctx: &Context,
1399 _operands: &[Uniformity],
1400 ) -> Vec<Uniformity> {
1401 vec![Uniformity::Cube]
1402 }
1403}
1404
1405#[op_interface_impl]
1406impl CanonicalizeInterface for WhileOp {
1407 fn canonicalize(&self, ctx: &mut Context, _rewriter: &mut MatchRewriter) -> Result<()> {
1408 let results = self.get_operation().results(ctx);
1409 let before_block = self.before_block(ctx);
1410 let after_block = self.after_block(ctx);
1411 for (idx, res) in results.into_iter().enumerate().rev() {
1412 let before_carried = self.get_before_carried_value(ctx, idx);
1413 let after_carried = self.get_after_carried_value(ctx, idx);
1414 if !res.is_used(ctx)
1415 && only_used_for_forward::<ConditionOp>(ctx, before_block, idx + 1, before_carried)
1416 && only_used_for_forward::<YieldOp>(ctx, after_block, idx, after_carried)
1417 {
1418 remove_from_terminator::<ConditionOp>(ctx, before_block, idx + 1);
1419 remove_from_terminator::<YieldOp>(ctx, after_block, idx);
1420 self.remove_initial_carried_value(ctx, idx);
1421 self.remove_before_carried_value(ctx, idx);
1422 self.remove_after_carried_value(ctx, idx);
1423 Operation::remove_result(self.get_operation(), ctx, idx);
1424 }
1425 }
1426 Ok(())
1427 }
1428}
1429
1430#[op_interface]
1431trait BranchToSCFOp {
1432 verify_op_succ!();
1433 fn to_scf(
1434 &self,
1435 ctx: &mut Context,
1436 rewriter: &mut DialectConversionRewriter,
1437 operands_info: &OperandsInfo,
1438 ) -> Result<()>;
1439}
1440
1441pub type BranchToSCFPass = DialectConversionPass<BranchToSCFConversion>;
1442
1443#[derive(Default, NamedRewrite)]
1444pub struct BranchToSCFConversion;
1445
1446impl DialectConversion for BranchToSCFConversion {
1447 fn can_convert_op(&self, ctx: &Context, op: Ptr<Operation>) -> bool {
1448 op.impls::<dyn BranchToSCFOp>(ctx)
1449 }
1450
1451 fn rewrite(
1452 &mut self,
1453 ctx: &mut Context,
1454 rewriter: &mut DialectConversionRewriter,
1455 op: Ptr<Operation>,
1456 operands_info: &OperandsInfo,
1457 ) -> Result<()> {
1458 let dyn_op = op.dyn_op(ctx);
1459 let to_scf = op_cast::<dyn BranchToSCFOp>(&*dyn_op).unwrap();
1460 to_scf.to_scf(ctx, rewriter, operands_info)
1461 }
1462}
1463
1464fn update_terminator<T: Op>(
1465 ctx: &Context,
1466 block: Ptr<BasicBlock>,
1467 default_reaching_def: Value,
1468 reaching_at_block_end: &HMap<Ptr<BasicBlock>, Value>,
1469) {
1470 let Some(term) = block.deref(ctx).get_terminator(ctx) else {
1471 return;
1472 };
1473 if !term.is_op::<T>(ctx) {
1474 return;
1475 }
1476 let block_reaching_def = reaching_at_block_end.get(&block).copied();
1477 let block_reaching_def = block_reaching_def.unwrap_or(default_reaching_def);
1478 Operation::push_operand(term, ctx, block_reaching_def);
1479}
1480
1481fn remove_from_terminator<T: Op>(ctx: &Context, block: Ptr<BasicBlock>, idx: usize) {
1482 let Some(term) = block.deref(ctx).get_terminator(ctx) else {
1483 return;
1484 };
1485 if !term.is_op::<T>(ctx) {
1486 return;
1487 }
1488 Operation::remove_operand(term, ctx, idx);
1489}
1490
1491fn only_used_for_forward<T: Op>(
1494 ctx: &Context,
1495 block: Ptr<BasicBlock>,
1496 idx: usize,
1497 val: Value,
1498) -> bool {
1499 if !val.is_used(ctx) {
1500 return true;
1501 }
1502 let Some(term) = block.deref(ctx).get_terminator(ctx) else {
1503 return false;
1504 };
1505 if !term.is_op::<T>(ctx) {
1506 return false;
1507 }
1508 let uses = val.uses(ctx);
1509 if uses.len() > 1 {
1510 return false;
1511 }
1512 uses[0].user_op() == term && uses[0].find_index(ctx) == idx
1513}
1514
1515pub(crate) fn block_mem_val(
1516 block: Ptr<BasicBlock>,
1517 default_reaching_def: MemoryValue,
1518 reaching_at_block_end: &HMap<Ptr<BasicBlock>, MemoryValue>,
1519) -> (MemoryRegionPredecessor, MemoryValue) {
1520 let block_reaching_def = reaching_at_block_end.get(&block).copied();
1521 let block_reaching_def = block_reaching_def.unwrap_or(default_reaching_def);
1522 (MemoryRegionPredecessor::Block(block), block_reaching_def)
1523}