feat(branch_condition_sink): sink branch cond for better peephole for…

… slt+beq
JuniMay · Aug 12, 2024 · abc13d9 · abc13d9
1 parent 15a1b31
commit abc13d9
Show file tree

Hide file tree

Showing 3 changed files with 113 additions and 16 deletions.
diff --git a/src/bin/compiler.rs b/src/bin/compiler.rs
@@ -13,6 +13,7 @@ use orzcc::{
             adce::{Adce, ADCE},
             bool2cond::{Bool2Cond, BOOL2COND},
             branch2select::{Branch2Select, BRANCH2SELECT},
+            branch_condition_sink::{BranchConditionSink, BRANCH_CONDITION_SINK},
             constant_phi::{ElimConstantPhi, ELIM_CONSTANT_PHI},
             control_flow::{
                 BlockReorder,
@@ -111,6 +112,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
             pipe_basic.add_pass(INSTCOMBINE);
             pipe_basic.add_pass(SIMPLE_DCE);
             pipe_basic.add_pass(GCM);
+            pipe_basic.add_pass(BRANCH_CONDITION_SINK);
             pipe_basic.add_pass(GVN);
             pipe_basic.add_pass(CFG_SIMPLIFY);
             pipe_basic.add_pass(ELIM_CONSTANT_PHI);
@@ -151,6 +153,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
             passman.run_transform(ELIM_CONSTANT_PHI, &mut ir, 32);
             passman.run_transform(SIMPLE_DCE, &mut ir, 32);
             passman.run_transform(GCM, &mut ir, 32);
+            passman.run_transform(BRANCH_CONDITION_SINK, &mut ir, 1);
             passman.run_transform(INDVAR_SIMPLIFY, &mut ir, 1);
             passman.run_transform(CONSTANT_FOLDING, &mut ir, 32);
             passman.run_transform(CFG_SIMPLIFY, &mut ir, 32);
@@ -165,18 +168,19 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
             passman.run_transform(ELIM_CONSTANT_PHI, &mut ir, 32);
             passman.run_transform(SIMPLE_DCE, &mut ir, 32);
 
-            passman.run_transform(INLINE, &mut ir, 1);
-            passman.run_transform(CONSTANT_FOLDING, &mut ir, 32);
-            passman.run_transform(CFG_SIMPLIFY, &mut ir, 32);
-            passman.run_transform(SIMPLE_DCE, &mut ir, 32);
-            passman.run_transform(GLOBAL_DCE, &mut ir, 32);
+            // passman.run_transform(INLINE, &mut ir, 1);
+            // passman.run_transform(CONSTANT_FOLDING, &mut ir, 32);
+            // passman.run_transform(CFG_SIMPLIFY, &mut ir, 32);
+            // passman.run_transform(SIMPLE_DCE, &mut ir, 32);
+            // passman.run_transform(GLOBAL_DCE, &mut ir, 32);
 
-            // TODO: unroll earlier to combine load/store
-            passman.run_transform(LOOP_UNROLL, &mut ir, 2);
-            passman.run_transform(GCM, &mut ir, 32);
-            passman.run_transform(CONSTANT_FOLDING, &mut ir, 32);
-            passman.run_transform(CFG_SIMPLIFY, &mut ir, 32);
-            passman.run_transform(SIMPLE_DCE, &mut ir, 32);
+            // // TODO: unroll earlier to combine load/store
+            // passman.run_transform(LOOP_UNROLL, &mut ir, 2);
+            // passman.run_transform(GCM, &mut ir, 32);
+            // passman.run_transform(BRANCH_CONDITION_SINK, &mut ir, 1);
+            // passman.run_transform(CONSTANT_FOLDING, &mut ir, 32);
+            // passman.run_transform(CFG_SIMPLIFY, &mut ir, 32);
+            // passman.run_transform(SIMPLE_DCE, &mut ir, 32);
 
             passman.run_transform(LEGALIZE, &mut ir, 1);
 
@@ -186,11 +190,11 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
                 let iter = passman.run_pipeline(&mut ir, &pipe_basic, 32, 8);
                 println!("pipeline basic iterations: {}", iter);
 
-                let iter = passman.run_pipeline(&mut ir, &pipe_inline, 32, 8);
-                println!("pipeline inline iterations: {}", iter);
+                // let iter = passman.run_pipeline(&mut ir, &pipe_inline, 32, 8);
+                // println!("pipeline inline iterations: {}", iter);
 
-                let iter = passman.run_pipeline(&mut ir, &pipe_unroll, 1, 1);
-                println!("pipeline unroll iterations: {}", iter);
+                // let iter = passman.run_pipeline(&mut ir, &pipe_unroll, 1, 1);
+                // println!("pipeline unroll iterations: {}", iter);
 
                 // a little expensive, run once per round
                 passman.run_transform(ADCE, &mut ir, 1);
@@ -228,13 +232,16 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         if cmd.opt > 0 {
             riscv64::run_peephole(lower_ctx.mctx_mut(), &cmd.lower_cfg);
             SimplifyCfg::run(lower_ctx.mctx_mut(), &cmd.lower_cfg);
-            RegisterCoalescing::run::<RvLowerSpec>(&mut lower_ctx, &cmd.lower_cfg);
         }
 
         if let Some(emit_vcode) = &cmd.emit_vcode {
             std::fs::write(emit_vcode, format!("{}", lower_ctx.mctx().display()))?;
         }
 
+        if cmd.opt > 0 {
+            RegisterCoalescing::run::<RvLowerSpec>(&mut lower_ctx, &cmd.lower_cfg);
+        }
+
         lower_ctx.reg_alloc();
         lower_ctx.after_regalloc();
 
@@ -279,6 +286,7 @@ fn register_passes(passman: &mut PassManager) {
 
     GlobalValueNumbering::register(passman);
     Gcm::register(passman);
+    BranchConditionSink::register(passman);
 
     Legalize::register(passman);
     BlockReorder::register(passman);

diff --git a/src/ir/passes/branch_condition_sink.rs b/src/ir/passes/branch_condition_sink.rs
@@ -0,0 +1,88 @@
+//! Sink branch conditions prior to the branch instruction.
+//!
+//! ```orzir
+//! %0 = icmp.slt %a, %b: i1
+//! ...
+//! br %0, %then, %elses
+//! ```
+//! -->
+//!
+//! ```orzir
+//! %0 = icmp.slt %a, %b: i1
+//! br %0, %then, %else
+//! ```
+//!
+//! we need to make sure that the condition is only used by the branch
+//! instruction.
+
+use super::control_flow::CfgCanonicalize;
+use crate::{
+    collections::linked_list::{LinkedListContainerPtr, LinkedListNodePtr},
+    ir::{
+        passman::{GlobalPassMut, LocalPassMut, PassManager, PassResult, TransformPass},
+        Context,
+        Func,
+    },
+    utils::def_use::Usable,
+};
+
+pub const BRANCH_CONDITION_SINK: &str = "branch_condition_sink";
+
+pub struct BranchConditionSink;
+
+impl LocalPassMut for BranchConditionSink {
+    type Output = ();
+
+    fn run(&mut self, ctx: &mut Context, func: Func) -> PassResult<(Self::Output, bool)> {
+        let mut changed = false;
+
+        let mut cursor = func.cursor();
+
+        while let Some(block) = cursor.next(ctx) {
+            let tail = block.tail(ctx).unwrap();
+
+            if !tail.is_br(ctx) {
+                continue;
+            }
+
+            let cond = tail.operand(ctx, 0);
+
+            if cond.users(ctx).len() != 1 {
+                continue;
+            }
+
+            if let Some(cond_def) = cond.def_inst(ctx) {
+                // move the condition prior to the branch instruction
+                cond_def.unlink(ctx);
+                block.push_inst_before_terminator(ctx, cond_def);
+
+                changed = true;
+            }
+        }
+
+        Ok(((), changed))
+    }
+}
+
+impl GlobalPassMut for BranchConditionSink {
+    type Output = ();
+
+    fn run(&mut self, ctx: &mut Context) -> PassResult<(Self::Output, bool)> {
+        let mut changed = false;
+        for func in ctx.funcs() {
+            let (_, local_changed) = LocalPassMut::run(self, ctx, func)?;
+            changed |= local_changed;
+        }
+        Ok(((), changed))
+    }
+}
+
+impl TransformPass for BranchConditionSink {
+    fn register(passman: &mut PassManager) {
+        passman.register_transform(
+            BRANCH_CONDITION_SINK,
+            BranchConditionSink,
+            vec![Box::new(CfgCanonicalize)],
+        );
+    }
+}
diff --git a/src/ir/passes/mod.rs b/src/ir/passes/mod.rs
@@ -1,6 +1,7 @@
 pub mod adce;
 pub mod bool2cond;
 pub mod branch2select;
+pub mod branch_condition_sink;
 pub mod constant_phi;
 pub mod control_flow;
 pub mod fold;