WIP

Cargo.toml

@@ -12,7 +12,7 @@ repository = "https://github.com/bytecodealliance/regalloc2"
 
 [dependencies]
 log = { version = "0.4.8", default-features = false }
-smallvec = { version = "1.6.1", features = ["union"] }
+smallvec = { version = "1.6.1", features = ["union", "const_generics"] }
 rustc-hash = { version = "1.1.0", default-features = false }
 slice-group-by = { version = "0.3.0", default-features = false }
 hashbrown = "0.13.2"

src/ion/fast_alloc2.rs

@@ -626,6 +626,19 @@ impl<'a, F: Function> FastAlloc<'a, F> {
                     self.livein_locs
                         .push(LiveinLoc::init(data.preg(), data.stack_slot()));
                 }
+
+                // also save block params
+                for &vreg in self.func.block_params(block) {
+                    let data = &self.vregs[vreg.vreg()];
+                    trace!(
+                        " -> {} at {:?} and slot {:?} (param)",
+                        vreg,
+                        data.preg(),
+                        data.stack_slot()
+                    );
+                    self.livein_locs
+                        .push(LiveinLoc::init(data.preg(), data.stack_slot()));
+                }
             }
 
             // load livein locations if our predecessor has multiple successors
@@ -670,6 +683,51 @@ impl<'a, F: Function> FastAlloc<'a, F> {
                     loc_idx += 1;
                 }
                 self.liveins = liveins;
+
+                // also restore block params
+                for &vreg in self.func.block_params(block) {
+                    let data = &mut self.vregs[vreg.vreg()];
+                    let loc = self.livein_locs[loc_idx];
+                    trace!(
+                        " -> {} at {:?} and slot {:?} (param)",
+                        vreg,
+                        loc.preg(),
+                        loc.slot()
+                    );
+                    match (data.preg(), loc.preg()) {
+                        (Some(cur_preg), Some(loc_preg)) => {
+                            if cur_preg != loc_preg {
+                                self.clear_preg(cur_preg.index());
+                                self.clear_preg(loc_preg.index());
+                                self.assign_preg(
+                                    loc_preg,
+                                    VReg::new(vreg.vreg(), loc_preg.class()),
+                                );
+                            }
+                        }
+                        (Some(cur_preg), None) => {
+                            self.clear_preg(cur_preg.index());
+                        }
+                        (None, Some(loc_preg)) => {
+                            self.clear_preg(loc_preg.index());
+                            self.assign_preg(loc_preg, VReg::new(vreg.vreg(), loc_preg.class()));
+                        }
+                        (None, None) => {}
+                    }
+                    let data = &mut self.vregs[vreg.vreg()];
+                    match loc.slot() {
+                        Some(slot) => {
+                            debug_assert!(data.allocated_stack_slot().is_some());
+                            debug_assert_eq!(data.allocated_stack_slot().unwrap(), slot);
+                            data.set_stack_slot_valid();
+                        }
+                        None => {
+                            data.clear_stack_slot(true);
+                        }
+                    }
+
+                    loc_idx += 1;
+                }
             }
 
             self.alloc_block_insts(block)?;
@@ -1208,6 +1266,7 @@ impl<'a, F: Function> FastAlloc<'a, F> {
         return res;
     }
 
+    // TODO: check if the reftype count is correctly managed
     fn alloc_block_edge_single_unallocated(
         &mut self,
         block: Block,
@@ -1291,6 +1350,7 @@ impl<'a, F: Function> FastAlloc<'a, F> {
                         to: Allocation::stack(SpillSlot::new(slot as usize)),
                     },
                 ));
+                self.vregs[succ_param.vreg()].set_stack_slot_valid();
 
                 self.vregs[out_vreg.vreg()].cur_use_idx -= 1;
                 continue;
@@ -1349,12 +1409,233 @@ impl<'a, F: Function> FastAlloc<'a, F> {
         todo!()
     }
 
+    // TODO: check if the reftype count is correctly managed
     fn alloc_block_edge_multiple(
         &mut self,
         block: Block,
         succs: &[Block],
     ) -> Result<(), RegAllocError> {
-        todo!()
+        trace!("  -> Block has multiple unallocated edges. Copying allocations");
+        let mut slots_made_valid = SmallVec::<[VReg; 4]>::new();
+        let mut slots_made_invalid: SmallVec<[VReg; 4]> = SmallVec::<[VReg; 4]>::new();
+        let mut preg_allocs: SmallVec<[(u16, VReg); PReg::NUM_INDEX + 1]> = SmallVec::new();
+
+        for idx in 0..self.pregs.len() {
+            debug_assert!(idx < PReg::NUM_INDEX);
+            let data = &self.pregs[idx];
+            match data.vreg() {
+                Some(vreg) => preg_allocs.push((idx as u16, vreg)),
+                None => {}
+            }
+        }
+        preg_allocs.push((u16::MAX, VReg::invalid()));
+
+        let last_inst = self.func.block_insns(block).last();
+        debug_assert!(self.func.is_branch(last_inst));
+        let out_vregs = self.func.branch_blockparams(block, last_inst, 0);
+
+        // we basically run the block_edge_single_unallocated algo but inside the target block and rewind afterwards
+        for &succ in succs {
+            assert_ne!(block, succ);
+            trace!("   -> Allocating edge for block {:?}", succ);
+
+            let succ_params = self.func.block_params(succ);
+            // prevent these from being allocated as tmp regs
+            let mut pregs_allocated = PRegSet::empty();
+            let mut tmp_reg = None;
+            for i in 0..succ_params.len() {
+                let succ_param = succ_params[i];
+                let out_vreg = out_vregs[i];
+
+                if self.vreg_killed(out_vreg.vreg(), block, self.cur_inst_pos + 1, false) {
+                    trace!(" -> {} mapping to {} is killed", out_vreg, succ_param);
+                    if let Some(slot) = self.vregs[out_vreg.vreg()].stack_slot() {
+                        trace!("   -> Transferring slot {} to {}", slot, succ_param);
+                        self.vregs[succ_param.vreg()].set_stack_slot_allocated(slot);
+                        self.vregs[succ_param.vreg()].set_stack_slot_valid();
+                        slots_made_valid.push(succ_param);
+                        if cfg!(debug_assertions) {
+                            self.vregs[out_vreg.vreg()].clear_stack_slot(false);
+                            slots_made_invalid.push(out_vreg);
+                        }
+                    }
+                    if let Some(preg) = self.vregs[out_vreg.vreg()].preg() {
+                        trace!("   -> Transferring {} to {}", preg, succ_param);
+                        self.vregs[succ_param.vreg()].set_preg(preg);
+                        if cfg!(debug_assertions) {
+                            self.vregs[out_vreg.vreg()].clear_preg();
+                        }
+                        self.pregs[preg.index()].set_vreg(succ_param);
+                        pregs_allocated.add(preg);
+                    }
+                    self.vregs[out_vreg.vreg()].cur_use_idx -= 1;
+                    continue;
+                }
+
+                trace!(" -> {} mapping to {} is not killed", out_vreg, succ_param);
+                // the out_vreg might only be alive until the end of the block but used by another param or still be live after the block
+                if let Some(preg) = self.vregs[out_vreg.vreg()].preg() {
+                    // TODO: we could do some fancy calculation here which of the block params should get to stay in the preg but rn we dont
+                    if self.next_use(out_vreg.vreg()).unwrap()
+                        > self.next_use(succ_param.vreg()).unwrap()
+                        || self.vregs[out_vreg.vreg()].stack_slot().is_some()
+                    {
+                        trace!("  -> Transferring preg from {} to {}", out_vreg, succ_param);
+
+                        if self.vregs[out_vreg.vreg()].stack_slot().is_none() {
+                            // save to stack if necessary
+                            self.alloc_and_move_to_stack(out_vreg, ProgPoint::before(last_inst));
+                            slots_made_valid.push(out_vreg);
+                        }
+
+                        self.vregs[out_vreg.vreg()].clear_preg();
+                        self.vregs[succ_param.vreg()].set_preg(preg);
+                        self.pregs[preg.index()].set_vreg(succ_param);
+                        pregs_allocated.add(preg);
+
+                        // don't need to care about stack slot
+                        self.vregs[out_vreg.vreg()].cur_use_idx -= 1;
+                        continue;
+                    }
+
+                    let slot = self.alloc_stack_slot(succ_param);
+                    trace!("  -> Transferring preg from {} to slot {}", preg, slot);
+                    self.edits.push((
+                        ProgPoint::before(last_inst),
+                        Edit::Move {
+                            from: Allocation::reg(preg),
+                            to: Allocation::stack(SpillSlot::new(slot as usize)),
+                        },
+                    ));
+                    self.vregs[succ_param.vreg()].set_stack_slot_valid();
+                    slots_made_valid.push(succ_param);
+
+                    self.vregs[out_vreg.vreg()].cur_use_idx -= 1;
+                    continue;
+                }
+
+                let dst_slot = self.alloc_stack_slot(succ_param) as usize;
+                let src_slot = self.vregs[out_vreg.vreg()].stack_slot().unwrap() as usize;
+                trace!(
+                    "  -> Copying {} from slot {} to slot {} for {}",
+                    out_vreg,
+                    src_slot,
+                    dst_slot,
+                    succ_param
+                );
+                // out_vreg is on the stack so create a stack slot for succ_param
+                let preg = match tmp_reg {
+                    Some(preg) => preg,
+                    None => {
+                        let preg = self.find_free_reg(
+                            succ_param.class(),
+                            pregs_allocated,
+                            Some(ProgPoint::before(last_inst)),
+                            block,
+                            self.cur_inst_pos + 1,
+                        )?;
+                        tmp_reg = Some(preg);
+                        preg
+                    }
+                };
+
+                self.edits.push((
+                    ProgPoint::before(last_inst),
+                    Edit::Move {
+                        from: Allocation::stack(SpillSlot::new(src_slot)),
+                        to: Allocation::reg(preg),
+                    },
+                ));
+                self.edits.push((
+                    ProgPoint::before(last_inst),
+                    Edit::Move {
+                        to: Allocation::stack(SpillSlot::new(dst_slot)),
+                        from: Allocation::reg(preg),
+                    },
+                ));
+                slots_made_valid.push(succ_param);
+                self.vregs[succ_param.vreg()].set_stack_slot_valid();
+
+                self.vregs[out_vreg.vreg()].cur_use_idx -= 1;
+            }
+
+            // save liveins
+            trace!(" -> Saving livein locations");
+            debug_assert_eq!(self.livein_loc_lookup[succ.index()], 0);
+            self.livein_loc_lookup[succ.index()] = self.livein_locs.len() as u32;
+            for vreg in self.liveins[succ.index()].iter() {
+                let data = &self.vregs[vreg];
+                trace!(
+                    "  -> {} at {:?} and slot {:?}",
+                    vreg,
+                    data.preg(),
+                    data.stack_slot()
+                );
+                self.livein_locs
+                    .push(LiveinLoc::init(data.preg(), data.stack_slot()));
+            }
+
+            // also save block params
+            for &vreg in self.func.block_params(succ) {
+                let data = &self.vregs[vreg.vreg()];
+                trace!(
+                    "  -> {} at {:?} and slot {:?} (param)",
+                    vreg,
+                    data.preg(),
+                    data.stack_slot()
+                );
+                self.livein_locs
+                    .push(LiveinLoc::init(data.preg(), data.stack_slot()));
+            }
+
+            // reset pregs & slots
+            trace!(" -> Resetting pregs and slots");
+            let mut vec_idx = 0;
+            for idx in 0..self.pregs.len() {
+                debug_assert!(idx < PReg::NUM_INDEX);
+                let data = &self.pregs[idx];
+                let next_preg = &preg_allocs[vec_idx];
+                if next_preg.0 as usize == idx {
+                    if data.vreg().is_none() || data.vreg().unwrap() != next_preg.1 {
+                        trace!("  -> Restoring {} to {}", idx, next_preg.1);
+                        self.clear_preg(idx);
+                        if let Some(preg) = self.vregs[next_preg.1.vreg()].preg() {
+                            self.clear_preg(preg.index());
+                        }
+                        let preg = PReg::new(idx, next_preg.1.class());
+                        self.assign_preg(preg, next_preg.1);
+                    }
+                    vec_idx += 1;
+                    continue;
+                }
+
+                self.clear_preg(idx);
+            }
+
+            for &vreg in &slots_made_valid {
+                trace!(
+                    "  -> Setting slot {} for {} valid",
+                    self.vregs[vreg.vreg()].allocated_stack_slot().unwrap(),
+                    vreg
+                );
+                self.vregs[vreg.vreg()].clear_stack_slot(true);
+            }
+            slots_made_valid.clear();
+
+            if cfg!(debug_assertions) {
+                for &vreg in &slots_made_invalid {
+                    trace!(
+                        "  -> Setting slot {} for {} invalid",
+                        self.vregs[vreg.vreg()].allocated_stack_slot().unwrap(),
+                        vreg
+                    );
+                    self.vregs[vreg.vreg()].set_stack_slot_valid();
+                }
+                slots_made_invalid.clear();
+            }
+        }
+
+        Ok(())
     }
 
     fn run(&mut self) -> Result<(), RegAllocError> {

src/lib.rs

@@ -1491,7 +1491,13 @@ pub fn run<F: Function>(
     env: &MachineEnv,
     options: &RegallocOptions,
 ) -> Result<Output, RegAllocError> {
-    ion::run(func, env, options.verbose_log, options.validate_ssa, options.fast_alloc)
+    ion::run(
+        func,
+        env,
+        options.verbose_log,
+        options.validate_ssa,
+        options.fast_alloc,
+    )
 }
 
 /// Options for allocation.