first steps to keeping stuff in register

src/ion/fast_alloc.rs

@@ -30,7 +30,7 @@ struct VRegData {
 
 #[derive(Default, Clone, Copy)]
 struct PRegData {
-    pub vreg: Option<u32>,
+    pub vreg: Option<VReg>,
     pub stack_pseudo: bool,
 }
 
@@ -153,6 +153,8 @@ struct FastAllocState<'a, F: Function> {
     pub stack_slot_count_int: u8,
     pub stack_slot_count_float: u8,
 
+    pub cur_inst_pos: usize,
+
     pub allocs: Vec<Allocation>,
     pub inst_alloc_offsets: Vec<u32>,
     pub edits: Vec<(ProgPoint, Edit)>,
@@ -227,6 +229,7 @@ impl<'a, F: Function> FastAllocState<'a, F> {
 
             cur_stack_slot_idx: 0,
             reftype_vregs_in_pregs_count: 0,
+            cur_inst_pos: 0,
 
             stack_slot_count_int: u8::try_from(func.spillslot_size(RegClass::Int))
                 .expect("that's a big integer"),
@@ -284,18 +287,30 @@ impl<'a, F: Function> FastAllocState<'a, F> {
 
     pub fn move_to_preg(&mut self, vreg: VReg, preg: PReg, pos: ProgPoint) {
         if let Some(vreg) = &self.pregs[preg.index()].vreg {
-            let vdata = &mut self.vregs[*vreg as usize];
+            let vdata = &mut self.vregs[vreg.vreg() as usize];
             debug_assert!(vdata.preg.is_some());
             debug_assert_eq!(vdata.preg.unwrap(), preg);
             vdata.preg = None;
         }
 
-        if let Some(preg) = &self.vregs[vreg.vreg()].preg {
+        if let Some(cur_preg) = &self.vregs[vreg.vreg()].preg {
+            // Do a reg->reg move
+            self.edits.push((
+                pos,
+                Edit::Move {
+                    from: Allocation::reg(*cur_preg),
+                    to: Allocation::reg(preg),
+                },
+            ));
+
             // TODO: allow multiple pregs for a single vreg?
-            let pdata = &mut self.pregs[preg.index()];
+            let pdata = &mut self.pregs[cur_preg.index()];
             debug_assert!(pdata.vreg.is_some());
-            debug_assert_eq!(pdata.vreg.unwrap(), vreg.vreg() as u32);
+            debug_assert_eq!(pdata.vreg.unwrap().vreg(), vreg.vreg());
             pdata.vreg = None;
+
+            self.pregs[preg.index()].vreg = Some(vreg);
+            return;
         }
 
         let vdata = &mut self.vregs[vreg.vreg()];
@@ -313,7 +328,7 @@ impl<'a, F: Function> FastAllocState<'a, F> {
             },
         ));
         vdata.preg = Some(preg);
-        pdata.vreg = Some(vreg.vreg() as u32);
+        pdata.vreg = Some(vreg);
 
         if vdata.reftype {
             self.reftype_vregs_in_pregs_count += 1;
@@ -327,7 +342,7 @@ impl<'a, F: Function> FastAllocState<'a, F> {
             panic!("Trying to move from unallocated preg/vreg to stack");
         }
         debug_assert_eq!(vdata.preg.unwrap(), preg);
-        debug_assert_eq!(pdata.vreg.unwrap(), vreg.vreg() as u32);
+        debug_assert_eq!(pdata.vreg.unwrap().vreg(), vreg.vreg());
 
         if vdata.slot_idx.is_none() {
             panic!("Trying to move to vreg without stack slot");
@@ -347,7 +362,7 @@ impl<'a, F: Function> FastAllocState<'a, F> {
         // need to make sure this is intended behavior
         self.clear_preg(preg);
 
-        self.pregs[preg.index()].vreg = Some(vreg.vreg() as u32);
+        self.pregs[preg.index()].vreg = Some(vreg);
         self.vregs[vreg.vreg()].preg = Some(preg);
         if self.vregs[vreg.vreg()].reftype {
             self.reftype_vregs_in_pregs_count += 1;
@@ -361,7 +376,7 @@ impl<'a, F: Function> FastAllocState<'a, F> {
     fn clear_preg_idx(&mut self, preg: usize) {
         let pdata = &mut self.pregs[preg];
         if let Some(vreg) = pdata.vreg {
-            let vdata = &mut self.vregs[vreg as usize];
+            let vdata = &mut self.vregs[vreg.vreg()];
             debug_assert_eq!(vdata.preg.unwrap().index(), preg);
             vdata.preg = None;
             pdata.vreg = None;
@@ -374,7 +389,7 @@ impl<'a, F: Function> FastAllocState<'a, F> {
     pub fn clear_vreg_from_reg(&mut self, vreg: VReg) {
         let vdata = &mut self.vregs[vreg.vreg()];
         if let Some(preg) = vdata.preg {
-            debug_assert_eq!(self.pregs[preg.index()].vreg.unwrap(), vreg.vreg() as u32);
+            debug_assert_eq!(self.pregs[preg.index()].vreg.unwrap().vreg(), vreg.vreg());
             self.pregs[preg.index()].vreg = None;
             vdata.preg = None;
 
@@ -391,7 +406,7 @@ impl<'a, F: Function> FastAllocState<'a, F> {
 
         for i in 0..self.pregs.len() {
             if let Some(vreg) = self.pregs[i].vreg.clone() {
-                let vreg = vreg as usize;
+                let vreg = vreg.vreg();
                 if self.vregs[vreg].reftype {
                     self.clear_preg_idx(i);
                 }
@@ -425,6 +440,11 @@ pub fn run<F: Function>(func: &F, mach_env: &MachineEnv) -> Result<Output, RegAl
 
         allocate_block_insts(&mut state, &const_state, block)?;
         handle_out_block_params(&mut state, &const_state, block)?;
+
+        let last_inst = state.func.block_insns(block).last();
+        if state.func.is_branch(last_inst) {
+            state.cur_inst_pos += 1;
+        }
     }
 
     // we do not iterate the blocks in their index order so the order of edits might not be sorted by progpoint
@@ -463,11 +483,39 @@ pub fn run<F: Function>(func: &F, mach_env: &MachineEnv) -> Result<Output, RegAl
     })
 }
 
+fn vreg_killed<'a, F: Function>(
+    state: &FastAllocState<'a, F>,
+    inst: Inst,
+    block: Block,
+    block_last_inst: usize,
+    vreg: usize,
+) -> bool {
+    let info = &state.vregs[vreg];
+    let block_after_pos = state.cur_inst_pos + (block_last_inst - inst.index()) + 1;
+    let cur_use_idx = info.cur_use_idx as usize;
+
+    if !state.liveouts[block.index()].get(vreg) {
+        if info.uses.len() >= cur_use_idx {
+            return true;
+        }
+        if info.uses[cur_use_idx] == state.cur_inst_pos as u32 {
+            if info.uses.len() >= cur_use_idx + 1
+                || info.uses[cur_use_idx + 1] >= block_after_pos as u32
+            {
+                return true;
+            }
+        }
+    }
+
+    return false;
+}
+
 fn allocate_block_insts<'a, F: Function>(
     state: &mut FastAllocState<'a, F>,
     const_state: &ReadOnlyData,
     block: Block,
 ) -> Result<(), RegAllocError> {
+    let block_last_inst_idx = state.func.block_insns(block).last().index();
     for inst in state.func.block_insns(block).iter() {
         let edit_start_idx = state.edits.len();
         let clobbers = state.func.inst_clobbers(inst);
@@ -537,6 +585,24 @@ fn allocate_block_insts<'a, F: Function>(
                 // TODO: treat them like normal vregs by just using last_vreg_index+1 for them?
                 match op.constraint() {
                     OperandConstraint::FixedReg(reg) => {
+                        // Save vreg if needed
+                        if let Some(vreg) = state.pregs[reg.index()].vreg {
+                            let vreg = vreg.vreg();
+                            if state.vregs[vreg].slot_idx.is_none()
+                                && !vreg_killed(state, inst, block, block_last_inst_idx, vreg)
+                            {
+                                let slot = state.create_stack_slot(reg.class());
+                                state.vregs[vreg].slot_idx = Some(slot);
+                                state.edits.push((
+                                    ProgPoint::before(inst),
+                                    Edit::Move {
+                                        from: Allocation::reg(reg),
+                                        to: Allocation::stack(SpillSlot::new(slot as usize)),
+                                    },
+                                ));
+                            }
+                        }
+
                         state.clear_preg(reg);
                         regs_allocated.add(reg);
                         state.allocs[alloc_idx + i] = Allocation::reg(reg);
@@ -562,7 +628,8 @@ fn allocate_block_insts<'a, F: Function>(
                                 return Err(RegAllocError::TooManyLiveRegs);
                             }
 
-                            if state.pregs[reg.index()].vreg.is_some() {
+                            // TODO: make this proper
+                            if regs_allocated.contains(reg) {
                                 // if the reg was allocated by another early use/write or late use
                                 // OR it is allocated and we have a late use we cannot do a correct allocation
                                 if op.pos() == OperandPos::Late || !late_write_regs.contains(reg) {
@@ -571,7 +638,48 @@ fn allocate_block_insts<'a, F: Function>(
                                 }
                             }
 
-                            state.move_to_preg(vreg, reg, ProgPoint::before(inst));
+                            // are we already in the correct reg?
+                            if let Some(cur_preg) = state.vregs[vreg.vreg()].preg {
+                                if cur_preg == reg {
+                                    trace!("{} already in target {}", vreg, cur_preg);
+                                    continue;
+                                }
+                            }
+
+                            // Save vreg if needed
+                            if let Some(vreg) = state.pregs[reg.index()].vreg {
+                                let vreg = vreg.vreg();
+                                if state.vregs[vreg].slot_idx.is_none()
+                                    && !vreg_killed(state, inst, block, block_last_inst_idx, vreg)
+                                {
+                                    let slot = state.create_stack_slot(reg.class());
+                                    state.vregs[vreg].slot_idx = Some(slot);
+                                    state.edits.push((
+                                        ProgPoint::before(inst),
+                                        Edit::Move {
+                                            from: Allocation::reg(reg),
+                                            to: Allocation::stack(SpillSlot::new(slot as usize)),
+                                        },
+                                    ));
+                                }
+                            }
+
+                            if let Some(cur_preg) = state.vregs[vreg.vreg()].preg {
+                                // Move from preg to preg
+                                state.edits.push((
+                                    ProgPoint::before(inst),
+                                    Edit::Move {
+                                        from: Allocation::reg(cur_preg),
+                                        to: Allocation::reg(reg),
+                                    },
+                                ));
+                                state.pregs[cur_preg.index()].vreg = None;
+                                state.vregs[vreg.vreg()].preg = Some(reg);
+                                state.pregs[reg.index()].vreg = Some(vreg);
+                            } else {
+                                state.move_to_preg(vreg, reg, ProgPoint::before(inst));
+                            }
+
                             state.allocs[alloc_idx + i] = Allocation::reg(reg);
                             if op.pos() == OperandPos::Late {
                                 if clobbers.contains(reg) {
@@ -605,6 +713,24 @@ fn allocate_block_insts<'a, F: Function>(
                                 late_write_disallow_regs.add(reg);
                             }
 
+                            // Save vreg if needed
+                            if let Some(vreg) = state.pregs[reg.index()].vreg {
+                                let vreg = vreg.vreg();
+                                if state.vregs[vreg].slot_idx.is_none()
+                                    && !vreg_killed(state, inst, block, block_last_inst_idx, vreg)
+                                {
+                                    let slot = state.create_stack_slot(reg.class());
+                                    state.vregs[vreg].slot_idx = Some(slot);
+                                    state.edits.push((
+                                        ProgPoint::before(inst),
+                                        Edit::Move {
+                                            from: Allocation::reg(reg),
+                                            to: Allocation::stack(SpillSlot::new(slot as usize)),
+                                        },
+                                    ));
+                                }
+                            }
+
                             state.vregs[vreg.vreg()].def_block = Some(block);
                             state.allocs[alloc_idx + i] = Allocation::reg(reg);
                             state.assign_preg(reg, vreg);
@@ -641,7 +767,30 @@ fn allocate_block_insts<'a, F: Function>(
 
                                 // Move from pseudoreg to tmp_reg and then to stack
                                 let tmp_reg = pregs.into_iter().next().unwrap();
-                                if state.pregs[tmp_reg.index()].vreg.is_some() {
+                                if let Some(vreg) = state.pregs[tmp_reg.index()].vreg {
+                                    // Save vreg if needed
+                                    let vreg = vreg.vreg();
+                                    if state.vregs[vreg].slot_idx.is_none()
+                                        && !vreg_killed(
+                                            state,
+                                            inst,
+                                            block,
+                                            block_last_inst_idx,
+                                            vreg,
+                                        )
+                                    {
+                                        let slot = state.create_stack_slot(reg.class());
+                                        state.vregs[vreg].slot_idx = Some(slot);
+                                        state.edits.push((
+                                            ProgPoint::before(inst),
+                                            Edit::Move {
+                                                from: Allocation::reg(tmp_reg),
+                                                to: Allocation::stack(SpillSlot::new(
+                                                    slot as usize,
+                                                )),
+                                            },
+                                        ));
+                                    }
                                     state.clear_preg(tmp_reg);
                                 }
 
@@ -716,6 +865,10 @@ fn allocate_block_insts<'a, F: Function>(
                             continue;
                         }
 
+                        if state.pregs[reg.index()].vreg.is_some() {
+                            continue;
+                        }
+
                         // reg should not contain anything
                         debug_assert!(state.pregs[reg.index()].vreg.is_none());
 
@@ -746,11 +899,74 @@ fn allocate_block_insts<'a, F: Function>(
                         continue;
                     }
 
-                    trace!("Ran out of registers for operand {}", i);
+                    trace!("No free register found for operand {}", i);
                     // No register available
-                    // TODO: try to evict vreg that does not need to be in a preg
+                    // TODO: first evict pregs that already have a stack slot even if they are used earlier?
-                    panic!("Out of registers: {:?}", regs_allocated);
+                    let mut evict_candidate = None;
-                    return Err(RegAllocError::TooManyLiveRegs);
+                    for &reg in reg_order {
+                        if regs_allocated.contains(reg) {
+                            continue;
+                        }
+
+                        debug_assert!(state.pregs[reg.index()].vreg.is_some());
+                        let vreg = state.pregs[reg.index()].vreg.unwrap().vreg();
+                        let next_use =
+                            state.vregs[vreg].uses[state.vregs[vreg].cur_use_idx as usize];
+                        if next_use == state.cur_inst_pos as u32 {
+                            continue;
+                        }
+
+                        if let Some((_, pos)) = &evict_candidate {
+                            if *pos < next_use {
+                                evict_candidate = Some((reg, next_use));
+                            }
+                        } else {
+                            evict_candidate = Some((reg, next_use));
+                        }
+                    }
+
+                    if let Some((reg, next_use)) = evict_candidate {
+                        // Save vreg if needed
+                        let vreg = state.pregs[reg.index()].vreg.unwrap();
+                        log::trace!("Evicting {} with v{}", reg, vreg);
+                        if state.vregs[vreg.vreg()].slot_idx.is_none()
+                            && !vreg_killed(state, inst, block, block_last_inst_idx, vreg.vreg())
+                        {
+                            let slot = state.create_stack_slot(reg.class());
+                            state.vregs[vreg.vreg()].slot_idx = Some(slot);
+                            state.edits.push((
+                                ProgPoint::before(inst),
+                                Edit::Move {
+                                    from: Allocation::reg(reg),
+                                    to: Allocation::stack(SpillSlot::new(slot as usize)),
+                                },
+                            ));
+                        }
+                        state.clear_preg(reg);
+
+                        state.allocs[alloc_idx + i] = Allocation::reg(reg);
+                        regs_allocated.add(reg);
+                        if op.kind() == OperandKind::Use {
+                            if req_refs_on_stack && state.vregs[vreg.vreg()].reftype {
+                                panic!("reftype required to be in reg at safepoint");
+                                return Err(RegAllocError::TooManyLiveRegs);
+                            }
+
+                            // need to move from stack to reg
+                            state.move_to_preg(vreg, reg, ProgPoint::before(inst));
+                        } else {
+                            // early def
+                            state.vregs[vreg.vreg()].def_block = Some(block);
+                            state.assign_preg(reg, vreg);
+                            state.alloc_stack_slot(vreg);
+                            state.move_to_stack(reg, vreg, ProgPoint::after(inst));
+                        }
+
+                        trace!("Chose {} for operand {}", reg, i);
+                    } else {
+                        panic!("Out of registers: {:?}", regs_allocated);
+                        return Err(RegAllocError::TooManyLiveRegs);
+                    }
                 }
                 OperandConstraint::Reuse(_) => {
                     panic!("Illegal register constraint reuse for early def or use");
@@ -783,6 +999,9 @@ fn allocate_block_insts<'a, F: Function>(
                         if regs_allocated.contains(reg) || late_write_disallow_regs.contains(reg) {
                             continue;
                         }
+                        if state.pregs[reg.index()].vreg.is_some() {
+                            continue;
+                        }
 
                         // reg should not contain anything
                         regs_allocated.add(reg);
@@ -801,10 +1020,62 @@ fn allocate_block_insts<'a, F: Function>(
                         continue;
                     }
 
-                    // No register available
+                    trace!("No free register found for {}", vreg);
-                    // TODO: try to evict vreg that does not need to be in a preg
+
-                    panic!("out of registers");
+                    // TODO: first evict pregs that already have a stack slot even if they are used earlier?
-                    return Err(RegAllocError::TooManyLiveRegs);
+                    let mut evict_candidate = None;
+                    for &reg in reg_order {
+                        if regs_allocated.contains(reg) || late_write_disallow_regs.contains(reg) {
+                            continue;
+                        }
+
+                        debug_assert!(state.pregs[reg.index()].vreg.is_some());
+                        let vreg = state.pregs[reg.index()].vreg.unwrap().vreg();
+                        let next_use =
+                            state.vregs[vreg].uses[state.vregs[vreg].cur_use_idx as usize];
+                        if next_use == state.cur_inst_pos as u32 {
+                            continue;
+                        }
+
+                        if let Some((_, pos)) = &evict_candidate {
+                            if *pos < next_use {
+                                evict_candidate = Some((reg, next_use));
+                            }
+                        } else {
+                            evict_candidate = Some((reg, next_use));
+                        }
+                    }
+
+                    if let Some((reg, next_use)) = evict_candidate {
+                        // Save vreg if needed
+                        let vreg = state.pregs[reg.index()].vreg.unwrap();
+                        log::trace!("Evicting {} with v{}", reg, vreg);
+                        if state.vregs[vreg.vreg()].slot_idx.is_none()
+                            && !vreg_killed(state, inst, block, block_last_inst_idx, vreg.vreg())
+                        {
+                            let slot = state.create_stack_slot(reg.class());
+                            state.vregs[vreg.vreg()].slot_idx = Some(slot);
+                            state.edits.push((
+                                ProgPoint::before(inst),
+                                Edit::Move {
+                                    from: Allocation::reg(reg),
+                                    to: Allocation::stack(SpillSlot::new(slot as usize)),
+                                },
+                            ));
+                        }
+                        state.clear_preg(reg);
+
+                        regs_allocated.add(reg);
+                        state.allocs[alloc_idx + i] = Allocation::reg(reg);
+
+                        state.assign_preg(reg, vreg);
+                        state.alloc_stack_slot(vreg);
+                        state.move_to_stack(reg, vreg, ProgPoint::after(inst));
+                        trace!("Chose {} for operand {}", reg, i);
+                    } else {
+                        panic!("out of registers");
+                        return Err(RegAllocError::TooManyLiveRegs);
+                    }
                 }
                 OperandConstraint::Reuse(idx) => {
                     debug_assert!(state.allocs[alloc_idx + idx].is_reg());
@@ -812,10 +1083,28 @@ fn allocate_block_insts<'a, F: Function>(
                     debug_assert!(regs_allocated.contains(preg));
                     state.allocs[alloc_idx + i] = Allocation::reg(preg);
 
+                    // Save vreg on stack if it is not killed
+                    if let Some(vreg) = state.pregs[preg.index()].vreg {
+                        let vreg = vreg.vreg();
+                        if state.vregs[vreg].slot_idx.is_none()
+                            && !vreg_killed(state, inst, block, block_last_inst_idx, vreg)
+                        {
+                            let slot = state.create_stack_slot(preg.class());
+                            state.vregs[vreg].slot_idx = Some(slot);
+                            state.edits.push((
+                                ProgPoint::before(inst),
+                                Edit::Move {
+                                    from: Allocation::reg(preg),
+                                    to: Allocation::stack(SpillSlot::new(slot as usize)),
+                                },
+                            ));
+                        }
+                    }
+
                     state.clear_preg(preg);
                     state.assign_preg(preg, vreg);
-                    state.alloc_stack_slot(vreg);
+                    //state.alloc_stack_slot(vreg);
-                    state.move_to_stack(preg, vreg, ProgPoint::after(inst));
+                    //state.move_to_stack(preg, vreg, ProgPoint::after(inst));
                 }
                 _ => {
                     debug_assert!(!state.allocs[alloc_idx + i].is_none());
@@ -823,10 +1112,24 @@ fn allocate_block_insts<'a, F: Function>(
             }
         }
 
-        // clear out all allocated regs
+        for op in operands {
-        for reg in regs_allocated {
+            if op.kind() != OperandKind::Use || op.as_fixed_nonallocatable().is_some() {
-            trace!("Clearing {}", reg);
+                continue;
-            state.clear_preg(reg);
+            }
+
+            let vreg_idx = op.vreg().vreg();
+            let info = &mut state.vregs[vreg_idx];
+            info.cur_use_idx += 1;
+
+            let block_after_pos = state.cur_inst_pos + (block_last_inst_idx - inst.index()) + 1;
+            // check if vreg dies here
+            if !state.liveouts[block.index()].get(vreg_idx)
+                && (info.uses.len() >= info.cur_use_idx as usize
+                    || info.uses[info.cur_use_idx as usize] > block_after_pos as u32)
+            {
+                // TODO: clear stack slot
+                state.clear_vreg_from_reg(op.vreg());
+            }
         }
 
         // fixup edit order
@@ -847,6 +1150,33 @@ fn allocate_block_insts<'a, F: Function>(
                 }
             }
         }
+
+        state.cur_inst_pos += 1;
+    }
+
+    // Move all liveout vregs to a stack slot if they dont have one and clear pregs
+    for i in 0..state.pregs.len() {
+        match state.pregs[i].vreg {
+            None => {}
+            Some(vreg) => {
+                let idx = vreg.vreg();
+                if state.liveouts[block.index()].get(idx) && state.vregs[idx].slot_idx.is_none() {
+                    let preg = state.vregs[idx].preg.unwrap();
+                    let slot = state.create_stack_slot(preg.class());
+                    state.edits.push((
+                        ProgPoint::before(Inst::new(block_last_inst_idx)),
+                        Edit::Move {
+                            from: Allocation::reg(preg),
+                            to: Allocation::stack(SpillSlot::new(slot as usize)),
+                        },
+                    ));
+                    state.vregs[idx].slot_idx = Some(slot);
+                }
+
+                state.vregs[idx].preg = None;
+                state.pregs[i].vreg = None;
+            }
+        }
     }
 
     Ok(())