[WIP] Add stack2reg optimization pass

7 files changed, 256 insertions, 22 deletions

diff --git a/example/mini_core.rs b/example/mini_core.rs
index 1912cc50a21..f649c945c05 100644
--- a/example/mini_core.rs
+++ b/example/mini_core.rs
@@ -394,6 +394,10 @@ pub trait FnMut<Args>: FnOnce<Args> {
 
 #[lang = "panic"]
 pub fn panic(&(_msg, _file, _line, _col): &(&'static str, &'static str, u32, u32)) -> ! {
+    panic_inner(&_msg);
+}
+
+pub fn panic_inner(_msg: &&str) -> ! {
     unsafe {
         libc::puts("Panicking\0" as *const str as *const u8);
         intrinsics::abort();
diff --git a/src/base.rs b/src/base.rs
index 6fc4f2ecddb..a5923114109 100644
--- a/src/base.rs
+++ b/src/base.rs
@@ -66,7 +66,7 @@ pub fn trans_fn<'clif, 'tcx, B: Backend + 'static>(
 
     // Recover all necessary data from fx, before accessing func will prevent future access to it.
     let instance = fx.instance;
-    let clif_comments = fx.clif_comments;
+    let mut clif_comments = fx.clif_comments;
     let source_info_set = fx.source_info_set;
     let local_map = fx.local_map;
 
@@ -76,6 +76,8 @@ pub fn trans_fn<'clif, 'tcx, B: Backend + 'static>(
     // Verify function
     verify_func(tcx, &clif_comments, &func);
 
+    crate::optimize::optimize_function(cx.tcx, instance, &mut func, &mut clif_comments);
+
     // Define function
     let context = &mut cx.caches.context;
     context.func = func;
@@ -108,7 +110,7 @@ pub fn trans_fn<'clif, 'tcx, B: Backend + 'static>(
     context.clear();
 }
 
-fn verify_func(tcx: TyCtxt, writer: &crate::pretty_clif::CommentWriter, func: &Function) {
+pub fn verify_func(tcx: TyCtxt, writer: &crate::pretty_clif::CommentWriter, func: &Function) {
     let flags = settings::Flags::new(settings::builder());
     match ::cranelift_codegen::verify_function(&func, &flags) {
         Ok(_) => {}
@@ -265,7 +267,7 @@ fn trans_stmt<'tcx>(
 
     fx.set_debug_loc(stmt.source_info);
 
-    #[cfg(debug_assertions)]
+    #[cfg(false_debug_assertions)]
     match &stmt.kind {
         StatementKind::StorageLive(..) | StatementKind::StorageDead(..) => {} // Those are not very useful
         _ => {
diff --git a/src/lib.rs b/src/lib.rs
index b01659eeb36..02cfe160370 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -47,6 +47,7 @@ mod linkage;
 mod main_shim;
 mod metadata;
 mod num;
+mod optimize;
 mod pointer;
 mod pretty_clif;
 mod target_features_whitelist;
diff --git a/src/optimize/mod.rs b/src/optimize/mod.rs
new file mode 100644
index 00000000000..f1ced869d26
--- /dev/null
+++ b/src/optimize/mod.rs
@@ -0,0 +1,15 @@
+use crate::prelude::*;
+
+mod stack2reg;
+
+pub fn optimize_function<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    instance: Instance<'tcx>,
+    func: &mut Function,
+    clif_comments: &mut crate::pretty_clif::CommentWriter,
+) {
+    self::stack2reg::optimize_function(func, clif_comments, format!("{:?}", instance));
+    #[cfg(debug_assertions)]
+    crate::pretty_clif::write_clif_file(tcx, "stack2reg", instance, &*func, &*clif_comments, None);
+    crate::base::verify_func(tcx, &*clif_comments, &*func);
+}
diff --git a/src/optimize/stack2reg.rs b/src/optimize/stack2reg.rs
new file mode 100644
index 00000000000..1737696c13e
--- /dev/null
+++ b/src/optimize/stack2reg.rs
@@ -0,0 +1,222 @@
+use cranelift_codegen::cursor::{Cursor, FuncCursor};
+use cranelift_codegen::ir::{Opcode, InstructionData, ValueDef};
+use cranelift_codegen::ir::immediates::Offset32;
+use cranelift_codegen::entity::SecondaryMap;
+
+use crate::prelude::*;
+
+pub(super) fn optimize_function(
+    func: &mut Function,
+    clif_comments: &mut crate::pretty_clif::CommentWriter,
+    name: String, // FIXME remove
+) {
+    let mut stack_addr_insts = SecondaryMap::new();
+    let mut stack_load_store_insts = SecondaryMap::new();
+
+    let mut cursor = FuncCursor::new(func);
+    while let Some(_ebb) = cursor.next_ebb() {
+        while let Some(inst) = cursor.next_inst() {
+            match cursor.func.dfg[inst] {
+                // Record all stack_addr, stack_load and stack_store instructions.
+                InstructionData::StackLoad {
+                    opcode: Opcode::StackAddr,
+                    stack_slot: _,
+                    offset: _,
+                } => {
+                    stack_addr_insts[inst] = true;
+                }
+                InstructionData::StackLoad {
+                    opcode: Opcode::StackLoad,
+                    stack_slot: _,
+                    offset: _,
+                } => {
+                    stack_load_store_insts[inst] = true;
+                }
+                InstructionData::StackStore {
+                    opcode: Opcode::StackStore,
+                    arg: _,
+                    stack_slot: _,
+                    offset: _,
+                } => {
+                    stack_load_store_insts[inst] = true;
+                }
+
+                // Turn load and store into stack_load and stack_store when possible.
+                InstructionData::Load { opcode: Opcode::Load, arg: addr, flags: _, offset } => {
+                    if cursor.func.dfg.ctrl_typevar(inst) == types::I128 || cursor.func.dfg.ctrl_typevar(inst).is_vector() {
+                        continue; // WORKAROUD: stack_load.i128 not yet implemented
+                    }
+                    if let Some((stack_slot, stack_addr_offset)) = try_get_stack_slot_and_offset_for_addr(cursor.func, addr) {
+                        if let Some(combined_offset) = offset.try_add_i64(stack_addr_offset.into()) {
+                            let ty = cursor.func.dfg.ctrl_typevar(inst);
+                            cursor.func.dfg.replace(inst).stack_load(ty, stack_slot, combined_offset);
+                            stack_load_store_insts[inst] = true;
+                        }
+                    }
+                }
+                InstructionData::Store { opcode: Opcode::Store, args: [value, addr], flags: _, offset } => {
+                    if cursor.func.dfg.ctrl_typevar(inst) == types::I128 || cursor.func.dfg.ctrl_typevar(inst).is_vector() {
+                        continue; // WORKAROUND: stack_store.i128 not yet implemented
+                    }
+                    if let Some((stack_slot, stack_addr_offset)) = try_get_stack_slot_and_offset_for_addr(cursor.func, addr) {
+                        if let Some(combined_offset) = offset.try_add_i64(stack_addr_offset.into()) {
+                            cursor.func.dfg.replace(inst).stack_store(value, stack_slot, combined_offset);
+                            stack_load_store_insts[inst] = true;
+                        }
+                    }
+                }
+                _ => {}
+            }
+        }
+    }
+
+    let mut used_stack_addr_insts = SecondaryMap::new();
+
+    let mut cursor = FuncCursor::new(func);
+    while let Some(_ebb) = cursor.next_ebb() {
+        while let Some(inst) = cursor.next_inst() {
+            for &arg in cursor.func.dfg.inst_args(inst) {
+                if let ValueDef::Result(arg_origin, 0) = cursor.func.dfg.value_def(arg) {
+                    if cursor.func.dfg[arg_origin].opcode() == Opcode::StackAddr {
+                        used_stack_addr_insts[arg_origin] = true;
+                    }
+                }
+            }
+        }
+    }
+
+    /*println!(
+        "stack_addr: [{}] ([{}] used)\nstack_load/stack_store: [{}]",
+        bool_secondary_map_to_string(&stack_addr_insts),
+        bool_secondary_map_to_string(&used_stack_addr_insts),
+        bool_secondary_map_to_string(&stack_load_store_insts),
+    );*/
+
+    for inst in used_stack_addr_insts.keys().filter(|&inst| used_stack_addr_insts[inst]) {
+        assert!(stack_addr_insts[inst]);
+    }
+
+    // Replace all unused stack_addr instructions with nop.
+    for inst in stack_addr_insts.keys() {
+        if stack_addr_insts[inst] && !used_stack_addr_insts[inst] {
+            func.dfg.detach_results(inst);
+            func.dfg.replace(inst).nop();
+            stack_addr_insts[inst] = false;
+        }
+    }
+
+    //println!("stack_addr (after): [{}]", bool_secondary_map_to_string(&stack_addr_insts));
+
+    let mut stack_slot_usage_map: SecondaryMap<StackSlot, HashSet<Inst>> = SecondaryMap::new();
+    for inst in stack_load_store_insts.keys().filter(|&inst| stack_load_store_insts[inst]) {
+        match func.dfg[inst] {
+            InstructionData::StackLoad {
+                opcode: Opcode::StackLoad,
+                stack_slot,
+                offset: _,
+            } => {
+                stack_slot_usage_map[stack_slot].insert(inst);
+            }
+            InstructionData::StackStore {
+                opcode: Opcode::StackStore,
+                arg: _,
+                stack_slot,
+                offset: _,
+            } => {
+                stack_slot_usage_map[stack_slot].insert(inst);
+            }
+            ref data => unreachable!("{:?}", data),
+        }
+    }
+    for inst in stack_addr_insts.keys().filter(|&inst| stack_addr_insts[inst]) {
+        match func.dfg[inst] {
+            InstructionData::StackLoad {
+                opcode: Opcode::StackAddr,
+                stack_slot,
+                offset: _,
+            } => {
+                stack_slot_usage_map[stack_slot].insert(inst);
+            }
+            ref data => unreachable!("{:?}", data),
+        }
+    }
+
+    //println!("{:?}\n", stack_slot_usage_map);
+
+    for (stack_slot, users) in stack_slot_usage_map.iter_mut() {
+        let mut is_addr_leaked = false;
+        let mut is_loaded = false;
+        let mut is_stored = false;
+        for &user in users.iter() {
+            match func.dfg[user] {
+                InstructionData::StackLoad {
+                    opcode: Opcode::StackAddr,
+                    stack_slot,
+                    offset: _,
+                } => {
+                    is_addr_leaked = true;
+                }
+                InstructionData::StackLoad {
+                    opcode: Opcode::StackLoad,
+                    stack_slot,
+                    offset: _,
+                } => {
+                    is_loaded = true;
+                }
+                InstructionData::StackStore {
+                    opcode: Opcode::StackStore,
+                    arg: _,
+                    stack_slot,
+                    offset: _,
+                } => {
+                    is_stored = true;
+                }
+                ref data => unreachable!("{:?}", data),
+            }
+        }
+
+        if is_addr_leaked || (is_loaded && is_stored) {
+            continue;
+        }
+
+        if is_loaded {
+            println!("[{}] [BUG?] Reading uninitialized memory", name);
+        } else {
+            // Stored value never read; just remove reads.
+            for &user in users.iter() {
+                println!("[{}] Remove dead stack store {} of {}", name, user, stack_slot);
+                func.dfg.replace(user).nop();
+            }
+        }
+    }
+}
+
+fn try_get_stack_slot_and_offset_for_addr(func: &Function, addr: Value) -> Option<(StackSlot, Offset32)> {
+    if let ValueDef::Result(addr_inst, 0) = func.dfg.value_def(addr) {
+        if let InstructionData::StackLoad {
+            opcode: Opcode::StackAddr,
+            stack_slot,
+            offset,
+        } = func.dfg[addr_inst] {
+            return Some((stack_slot, offset));
+        }
+    }
+    None
+}
+
+fn bool_secondary_map_to_string<E>(map: &SecondaryMap<E, bool>) -> String
+    where E: cranelift_codegen::entity::EntityRef + std::fmt::Display,
+{
+    map
+        .keys()
+        .filter_map(|inst| {
+            // EntitySet::keys returns all possible entities until the last entity inserted.
+            if map[inst] {
+                Some(format!("{}", inst))
+            } else {
+                None
+            }
+        })
+        .collect::<Vec<String>>()
+        .join(", ")
+}
diff --git a/src/value_and_place.rs b/src/value_and_place.rs
index 39eb7d24707..6f0c615c950 100644
--- a/src/value_and_place.rs
+++ b/src/value_and_place.rs
@@ -362,19 +362,6 @@ impl<'tcx> CPlace<'tcx> {
     }
 
     pub fn write_cvalue(self, fx: &mut FunctionCx<'_, 'tcx, impl Backend>, from: CValue<'tcx>) {
-        #[cfg(debug_assertions)]
-        {
-            use cranelift_codegen::cursor::{Cursor, CursorPosition};
-            let cur_ebb = match fx.bcx.cursor().position() {
-                CursorPosition::After(ebb) => ebb,
-                _ => unreachable!(),
-            };
-            fx.add_comment(
-                fx.bcx.func.layout.last_inst(cur_ebb).unwrap(),
-                format!("write_cvalue: {:?} <- {:?}",self, from),
-            );
-        }
-
         let from_ty = from.layout().ty;
         let to_ty = self.layout().ty;
 
diff --git a/test.sh b/test.sh
index e502a20db5b..98171bc0f6a 100755
--- a/test.sh
+++ b/test.sh
@@ -4,10 +4,10 @@ set -e
 
 if [[ "$1" == "--release" ]]; then
     export CHANNEL='release'
-    cargo build --release $CG_CLIF_COMPILE_FLAGS
+    cargo rustc --release $CG_CLIF_COMPILE_FLAGS -- -Clink-args=-fuse-ld=lld
 else
     export CHANNEL='debug'
-    cargo build $CG_CLIF_COMPILE_FLAGS
+    cargo rustc $CG_CLIF_COMPILE_FLAGS -- -Clink-args=-fuse-ld=lld
 fi
 
 source config.sh
@@ -28,6 +28,8 @@ mkdir -p target/out/clif
 echo "[BUILD] mini_core"
 $RUSTC example/mini_core.rs --crate-name mini_core --crate-type lib,dylib
 
+#exit 1
+
 echo "[BUILD] example"
 $RUSTC example/example.rs --crate-type lib
 
@@ -69,15 +71,16 @@ $RUSTC example/mod_bench.rs --crate-type bin
 
 pushd simple-raytracer
 echo "[BENCH COMPILE] ebobby/simple-raytracer"
-hyperfine --runs ${RUN_RUNS:-10} --warmup 1 --prepare "rm -r target/*/debug || true" \
-    "RUSTFLAGS='' cargo build --target $TARGET_TRIPLE" \
-    "../cargo.sh build"
+rm -r target/x86_64*/
+../cargo.sh build
 
 echo "[BENCH RUN] ebobby/simple-raytracer"
 cp ./target/*/debug/main ./raytracer_cg_clif
-hyperfine --runs ${RUN_RUNS:-10} ./raytracer_cg_llvm ./raytracer_cg_clif
+hyperfine --runs ${RUN_RUNS:-10} ./raytracer_*
 popd
 
+exit 1
+
 pushd build_sysroot/sysroot_src/src/libcore/tests
 rm -r ./target || true
 ../../../../../cargo.sh test