Add IR printing functionality to the PassManager

Signed-off-by: Tomas Fabrizio Orsi <[email protected]>

Author: lima-limon-inc

dialects/hir/src/transforms/spill.rs

@@ -3,7 +3,7 @@ use alloc::rc::Rc;
 use midenc_hir::{
     adt::SmallDenseMap,
     dialects::builtin::{Function, FunctionRef, LocalVariable},
-    pass::{Pass, PassExecutionState},
+    pass::{Pass, PassExecutionState, PassIdentifier, PostPassStatus},
     BlockRef, BuilderExt, EntityMut, Op, OpBuilder, OperationName, OperationRef, Report, Rewriter,
     SourceSpan, Spanned, Symbol, ValueRef,
 };
@@ -19,6 +19,10 @@ impl Pass for TransformSpills {
         "transform-spills"
     }
 
+    fn pass_id(&self) -> Option<PassIdentifier> {
+        Some(PassIdentifier::TransformSpills)
+    }
+
     fn argument(&self) -> &'static str {
         "transform-spills"
     }
@@ -31,22 +35,22 @@ impl Pass for TransformSpills {
         &mut self,
         op: EntityMut<'_, Self::Target>,
         state: &mut PassExecutionState,
-    ) -> Result<(), Report> {
+    ) -> Result<PostPassStatus, Report> {
         let function = op.into_entity_ref();
         log::debug!(target: "insert-spills", "computing and inserting spills for {}", function.as_operation());
 
         if function.is_declaration() {
             log::debug!(target: "insert-spills", "function has no body, no spills needed!");
             state.preserved_analyses_mut().preserve_all();
-            return Ok(());
+            return Ok(PostPassStatus::IRUnchanged);
         }
         let mut analysis =
             state.analysis_manager().get_analysis_for::<SpillAnalysis, Function>()?;
 
         if !analysis.has_spills() {
             log::debug!(target: "insert-spills", "no spills needed!");
             state.preserved_analyses_mut().preserve_all();
-            return Ok(());
+            return Ok(PostPassStatus::IRUnchanged);
         }
 
         // Take ownership of the spills analysis so that we can mutate the analysis state during

dialects/scf/src/transforms/cfg_to_scf.rs

@@ -7,7 +7,7 @@ use midenc_hir::{
     diagnostics::Severity,
     dialects::builtin,
     dominance::DominanceInfo,
-    pass::{Pass, PassExecutionState},
+    pass::{Pass, PassExecutionState, PassIdentifier, PostPassStatus},
     Builder, EntityMut, Forward, Op, Operation, OperationName, OperationRef, RawWalk, Report,
     SmallVec, Spanned, Type, ValueRange, ValueRef, WalkResult,
 };
@@ -35,6 +35,10 @@ impl Pass for LiftControlFlowToSCF {
         "lift-control-flow"
     }
 
+    fn pass_id(&self) -> Option<PassIdentifier> {
+        Some(PassIdentifier::LiftControlFlowToSCF)
+    }
+
     fn argument(&self) -> &'static str {
         "lift-control-flow"
     }
@@ -51,7 +55,7 @@ impl Pass for LiftControlFlowToSCF {
         &mut self,
         op: EntityMut<'_, Self::Target>,
         state: &mut PassExecutionState,
-    ) -> Result<(), Report> {
+    ) -> Result<PostPassStatus, Report> {
         let mut transformation = ControlFlowToSCFTransformation;
         let mut changed = false;
 
@@ -130,7 +134,7 @@ impl Pass for LiftControlFlowToSCF {
         });
 
         if result.was_interrupted() {
-            return result.into_result();
+            return result.into_result().map(|_| PostPassStatus::IRUnchanged);
         }
 
         log::debug!(
@@ -141,7 +145,7 @@ impl Pass for LiftControlFlowToSCF {
             state.preserved_analyses_mut().preserve_all();
         }
 
-        Ok(())
+        Ok(changed.into())
     }
 }
 

hir-transform/src/canonicalization.rs

@@ -1,7 +1,7 @@
 use alloc::{boxed::Box, format, rc::Rc};
 
 use midenc_hir::{
-    pass::{OperationPass, Pass, PassExecutionState},
+    pass::{OperationPass, Pass, PassExecutionState, PassIdentifier, PostPassStatus},
     patterns::{self, FrozenRewritePatternSet, GreedyRewriteConfig, RewritePatternSet},
     Context, EntityMut, Operation, OperationName, Report, Spanned,
 };
@@ -62,6 +62,10 @@ impl Pass for Canonicalizer {
         "canonicalizer"
     }
 
+    fn pass_id(&self) -> Option<PassIdentifier> {
+        Some(PassIdentifier::Canonicalizer)
+    }
+
     fn argument(&self) -> &'static str {
         "canonicalizer"
     }
@@ -93,10 +97,10 @@ impl Pass for Canonicalizer {
         &mut self,
         op: EntityMut<'_, Self::Target>,
         state: &mut PassExecutionState,
-    ) -> Result<(), Report> {
+    ) -> Result<PostPassStatus, Report> {
         let Some(rewrites) = self.rewrites.as_ref() else {
             log::debug!("skipping canonicalization as there are no rewrite patterns to apply");
-            return Ok(());
+            return Ok(PostPassStatus::IRUnchanged);
         };
         let op = {
             let ptr = op.as_operation_ref();
@@ -129,16 +133,20 @@ impl Pass for Canonicalizer {
         }
 
         let op = op.borrow();
-        match converged {
+        let changed = match converged {
             Ok(changed) => {
-                log::debug!("canonicalization converged for '{}', changed={changed}", op.name())
+                log::debug!("canonicalization converged for '{}', changed={changed}", op.name());
+                changed
             }
-            Err(changed) => log::warn!(
-                "canonicalization failed to converge for '{}', changed={changed}",
-                op.name()
-            ),
-        }
+            Err(changed) => {
+                log::warn!(
+                    "canonicalization failed to converge for '{}', changed={changed}",
+                    op.name()
+                );
+                changed
+            }
+        };
 
-        Ok(())
+        Ok(changed.into())
     }
 }

hir-transform/src/sccp.rs

@@ -1,5 +1,5 @@
 use midenc_hir::{
-    pass::{Pass, PassExecutionState},
+    pass::{Pass, PassExecutionState, PassIdentifier, PostPassStatus},
     patterns::NoopRewriterListener,
     BlockRef, Builder, EntityMut, OpBuilder, Operation, OperationFolder, OperationName, RegionList,
     Report, SmallVec, ValueRef,
@@ -26,6 +26,10 @@ impl Pass for SparseConditionalConstantPropagation {
         "sparse-conditional-constant-propagation"
     }
 
+    fn pass_id(&self) -> Option<PassIdentifier> {
+        Some(PassIdentifier::SparseConditionalConstantPropagation)
+    }
+
     fn argument(&self) -> &'static str {
         "sparse-conditional-constant-propagation"
     }
@@ -38,7 +42,7 @@ impl Pass for SparseConditionalConstantPropagation {
         &mut self,
         mut op: EntityMut<'_, Self::Target>,
         state: &mut PassExecutionState,
-    ) -> Result<(), Report> {
+    ) -> Result<PostPassStatus, Report> {
         // Run sparse constant propagation + dead code analysis
         let mut solver = DataFlowSolver::default();
         solver.load::<DeadCodeAnalysis>();
@@ -58,7 +62,7 @@ impl SparseConditionalConstantPropagation {
         op: &mut Operation,
         _state: &mut PassExecutionState,
         solver: &DataFlowSolver,
-    ) -> Result<(), Report> {
+    ) -> Result<PostPassStatus, Report> {
         let mut worklist = SmallVec::<[BlockRef; 8]>::default();
 
         let add_to_worklist = |regions: &RegionList, worklist: &mut SmallVec<[BlockRef; 8]>| {
@@ -76,6 +80,7 @@ impl SparseConditionalConstantPropagation {
 
         add_to_worklist(op.regions(), &mut worklist);
 
+        let mut replaced_any = false;
         while let Some(mut block) = worklist.pop() {
             let mut block = block.borrow_mut();
             let body = block.body_mut();
@@ -91,8 +96,10 @@ impl SparseConditionalConstantPropagation {
                 let mut replaced_all = num_results != 0;
                 for index in 0..num_results {
                     let result = { op.borrow().get_result(index).borrow().as_value_ref() };
-                    replaced_all &=
-                        replace_with_constant(solver, &mut builder, &mut folder, result);
+                    let replaced = replace_with_constant(solver, &mut builder, &mut folder, result);
+
+                    replaced_any |= replaced;
+                    replaced_all &= replaced;
                 }
 
                 // If all of the results of the operation were replaced, try to erase the operation
@@ -112,7 +119,7 @@ impl SparseConditionalConstantPropagation {
             builder.set_insertion_point_to_start(block.as_block_ref());
 
             for arg in block.arguments() {
-                replace_with_constant(
+                replaced_any |= replace_with_constant(
                     solver,
                     &mut builder,
                     &mut folder,
@@ -121,7 +128,7 @@ impl SparseConditionalConstantPropagation {
             }
         }
 
-        Ok(())
+        Ok(replaced_any.into())
     }
 }
 

hir-transform/src/sink.rs

@@ -4,7 +4,7 @@ use midenc_hir::{
     adt::SmallDenseMap,
     dominance::DominanceInfo,
     matchers::{self, Matcher},
-    pass::{Pass, PassExecutionState},
+    pass::{Pass, PassExecutionState, PassIdentifier, PostPassStatus},
     traits::{ConstantLike, Terminator},
     Backward, Builder, EntityMut, Forward, FxHashSet, OpBuilder, Operation, OperationName,
     OperationRef, ProgramPoint, RawWalk, Region, RegionBranchOpInterface,
@@ -84,6 +84,10 @@ impl Pass for ControlFlowSink {
         "control-flow-sink"
     }
 
+    fn pass_id(&self) -> Option<PassIdentifier> {
+        Some(PassIdentifier::ControlFlowSink)
+    }
+
     fn argument(&self) -> &'static str {
         "control-flow-sink"
     }
@@ -96,7 +100,7 @@ impl Pass for ControlFlowSink {
         &mut self,
         op: EntityMut<'_, Self::Target>,
         state: &mut PassExecutionState,
-    ) -> Result<(), Report> {
+    ) -> Result<PostPassStatus, Report> {
         let op = op.into_entity_ref();
         log::debug!(target: "control-flow-sink", "sinking operations in {op}");
 
@@ -105,6 +109,7 @@ impl Pass for ControlFlowSink {
 
         let dominfo = state.analysis_manager().get_analysis::<DominanceInfo>()?;
 
+        let mut sunk = PostPassStatus::IRUnchanged;
         operation.raw_prewalk_all::<Forward, _>(|op: OperationRef| {
             let regions_to_sink = {
                 let op = op.borrow();
@@ -118,7 +123,7 @@ impl Pass for ControlFlowSink {
             };
 
             // Sink side-effect free operations.
-            control_flow_sink(
+            sunk = control_flow_sink(
                 &regions_to_sink,
                 &dominfo,
                 |op: &Operation, _region: &Region| op.is_memory_effect_free(),
@@ -132,7 +137,7 @@ impl Pass for ControlFlowSink {
             );
         });
 
-        Ok(())
+        Ok(sunk)
     }
 }
 
@@ -160,6 +165,10 @@ impl Pass for SinkOperandDefs {
         "sink-operand-defs"
     }
 
+    fn pass_id(&self) -> Option<PassIdentifier> {
+        Some(PassIdentifier::SinkOperandDefs)
+    }
+
     fn argument(&self) -> &'static str {
         "sink-operand-defs"
     }
@@ -172,7 +181,7 @@ impl Pass for SinkOperandDefs {
         &mut self,
         op: EntityMut<'_, Self::Target>,
         _state: &mut PassExecutionState,
-    ) -> Result<(), Report> {
+    ) -> Result<PostPassStatus, Report> {
         let operation = op.as_operation_ref();
         drop(op);
 
@@ -184,6 +193,7 @@ impl Pass for SinkOperandDefs {
         // then process the worklist, moving everything into position.
         let mut worklist = alloc::collections::VecDeque::default();
 
+        let mut changed = PostPassStatus::IRUnchanged;
         // Visit ops in "true" post-order (i.e. block bodies are visited bottom-up).
         operation.raw_postwalk_all::<Backward, _>(|operation: OperationRef| {
             // Determine if any of this operation's operands represent one of the following:
@@ -308,6 +318,7 @@ impl Pass for SinkOperandDefs {
                         log::trace!(target: "sink-operand-defs", "    rewriting operand {operand_value} as {replacement}");
                         operand.borrow_mut().set(replacement);
 
+                        changed = PostPassStatus::IRChanged;
                         // If no other uses of this value remain, then remove the original
                         // operation, as it is now dead.
                         if !operand_value.borrow().is_used() {
@@ -354,6 +365,7 @@ impl Pass for SinkOperandDefs {
                         log::trace!(target: "sink-operand-defs", "    rewriting operand {operand_value} as {replacement}");
                         sink_state.replacements.insert(operand_value, replacement);
                         operand.borrow_mut().set(replacement);
+                        changed = PostPassStatus::IRChanged;
                     } else {
                         log::trace!(target: "sink-operand-defs", "    defining op is a constant with no other uses, moving into place");
                         // The original op can be moved
@@ -397,7 +409,7 @@ impl Pass for SinkOperandDefs {
             }
         }
 
-        Ok(())
+        Ok(changed)
     }
 }
 
@@ -548,12 +560,14 @@ pub fn control_flow_sink<P, F>(
     dominfo: &DominanceInfo,
     should_move_into_region: P,
     move_into_region: F,
-) where
+) -> PostPassStatus
+where
     P: Fn(&Operation, &Region) -> bool,
     F: Fn(OperationRef, RegionRef),
 {
     let sinker = Sinker::new(dominfo, should_move_into_region, move_into_region);
-    sinker.sink_regions(regions);
+    let sunk_regions = sinker.sink_regions(regions);
+    (sunk_regions > 0).into()
 }
 
 /// Populates `regions` with regions of the provided region branch op that are executed at most once

hir-transform/src/spill.rs

@@ -4,7 +4,7 @@ use midenc_hir::{
     adt::{SmallDenseMap, SmallSet},
     cfg::Graph,
     dominance::{DomTreeNode, DominanceFrontier, DominanceInfo},
-    pass::AnalysisManager,
+    pass::{AnalysisManager, PostPassStatus},
     traits::SingleRegion,
     BlockRef, Builder, Context, FxHashMap, OpBuilder, OpOperand, Operation, OperationRef,
     ProgramPoint, Region, RegionBranchOpInterface, RegionBranchPoint, RegionRef, Report, Rewriter,
@@ -113,7 +113,7 @@ pub fn transform_spills(
     analysis: &mut SpillAnalysis,
     interface: &mut dyn TransformSpillsInterface,
     analysis_manager: AnalysisManager,
-) -> Result<(), Report> {
+) -> Result<PostPassStatus, Report> {
     assert!(
         op.borrow().implements::<dyn SingleRegion>(),
         "the spills transformation is not supported when the root op is multi-region"
@@ -252,7 +252,7 @@ pub fn transform_spills(
         )?;
     }
 
-    Ok(())
+    Ok(PostPassStatus::IRChanged)
 }
 
 fn rewrite_single_block_spills(

hir/src/ir/print.rs

@@ -9,6 +9,7 @@ use crate::{
     AttributeValue, EntityWithId, SuccessorOperands, Value,
 };
 
+#[derive(Debug)]
 pub struct OpPrintingFlags {
     pub print_entry_block_headers: bool,
     pub print_intrinsic_attributes: bool,