Loop depth instead of hot/cold, with fast O(n) loop-depth computation. Use this to compute use weights.

src/cfg.rs

@@ -6,6 +6,7 @@
 //! Lightweight CFG analyses.
 
 use crate::{domtree, postorder, Block, Function, Inst, OperandKind, ProgPoint, RegAllocError};
+use smallvec::{smallvec, SmallVec};
 
 #[derive(Clone, Debug)]
 pub struct CFGInfo {
@@ -34,6 +35,14 @@ pub struct CFGInfo {
     /// just one value per block and always know any block's position in its
     /// successors' preds lists.)
     pub pred_pos: Vec<usize>,
+    /// For each block, what is the approximate loop depth?
+    ///
+    /// This measure is fully precise iff the input CFG is reducible
+    /// and blocks are in RPO, so that loop backedges are precisely
+    /// those whose block target indices are less than their source
+    /// indices. Otherwise, it will be approximate, but should still
+    /// be usable for heuristic purposes.
+    pub approx_loop_depth: Vec<u32>,
 }
 
 impl CFGInfo {
@@ -52,6 +61,8 @@ impl CFGInfo {
         let mut block_entry = vec![ProgPoint::before(Inst::invalid()); f.blocks()];
         let mut block_exit = vec![ProgPoint::before(Inst::invalid()); f.blocks()];
         let mut pred_pos = vec![0; f.blocks()];
+        let mut backedge_in = vec![0; f.blocks()];
+        let mut backedge_out = vec![0; f.blocks()];
 
         for block in 0..f.blocks() {
             let block = Block::new(block);
@@ -104,6 +115,34 @@ impl CFGInfo {
                     return Err(RegAllocError::DisallowedBranchArg(last));
                 }
             }
+
+            for &succ in f.block_succs(block) {
+                if succ.index() <= block.index() {
+                    backedge_in[succ.index()] += 1;
+                    backedge_out[block.index()] += 1;
+                }
+            }
+        }
+
+        let mut approx_loop_depth = vec![];
+        let mut backedge_stack: SmallVec<[usize; 4]> = smallvec![];
+        let mut cur_depth = 0;
+        for block in 0..f.blocks() {
+            if backedge_in[block] > 0 {
+                cur_depth += 1;
+                backedge_stack.push(backedge_in[block]);
+            }
+
+            approx_loop_depth.push(cur_depth);
+
+            while backedge_stack.len() > 0 && backedge_out[block] > 0 {
+                backedge_out[block] -= 1;
+                *backedge_stack.last_mut().unwrap() -= 1;
+                if *backedge_stack.last().unwrap() == 0 {
+                    cur_depth -= 1;
+                    backedge_stack.pop();
+                }
+            }
         }
 
         Ok(CFGInfo {
@@ -115,6 +154,7 @@ impl CFGInfo {
             block_entry,
             block_exit,
             pred_pos,
+            approx_loop_depth,
         })
     }
 

src/ion/mod.rs

@@ -354,7 +354,6 @@ struct Env<'a, F: Function> {
     vreg_regs: Vec<VReg>,
     pregs: Vec<PRegData>,
     allocation_queue: PrioQueue,
-    hot_code: LiveRangeSet,
     clobbers: Vec<Inst>,   // Sorted list of insts with clobbers.
     safepoints: Vec<Inst>, // Sorted list of safepoint insts.
     safepoints_per_vreg: HashMap<usize, HashSet<Inst>>,
@@ -524,8 +523,10 @@ impl LiveRangeSet {
 }
 
 #[inline(always)]
-fn spill_weight_from_policy(policy: OperandPolicy, is_hot: bool, is_def: bool) -> u32 {
+fn spill_weight_from_policy(policy: OperandPolicy, loop_depth: usize, is_def: bool) -> u32 {
-    let hot_bonus = if is_hot { 10000 } else { 0 };
+    // A bonus of 1000 for one loop level, 4000 for two loop levels,
+    // 16000 for three loop levels, etc. Avoids exponentiation.
+    let hot_bonus = std::cmp::min(16000, 1000 * (1 << (2 * loop_depth)));
     let def_bonus = if is_def { 2000 } else { 0 };
     let policy_bonus = match policy {
         OperandPolicy::Any => 1000,
@@ -794,7 +795,6 @@ impl<'a, F: Function> Env<'a, F> {
             clobbers: vec![],
             safepoints: vec![],
             safepoints_per_vreg: HashMap::new(),
-            hot_code: LiveRangeSet::new(),
             spilled_bundles: vec![],
             spillslots: vec![],
             slots_by_size: vec![],
@@ -944,17 +944,12 @@ impl<'a, F: Function> Env<'a, F> {
     }
 
     fn insert_use_into_liverange(&mut self, into: LiveRangeIndex, mut u: Use) {
-        let insert_pos = u.pos;
         let operand = u.operand;
         let policy = operand.policy();
-        let is_hot = self
+        let block = self.cfginfo.insn_block[u.pos.inst().index()];
-            .hot_code
+        let loop_depth = self.cfginfo.approx_loop_depth[block.index()] as usize;
-            .btree
+        let weight =
-            .contains_key(&LiveRangeKey::from_range(&CodeRange {
+            spill_weight_from_policy(policy, loop_depth, operand.kind() != OperandKind::Use);
-                from: insert_pos,
-                to: insert_pos.next(),
-            }));
-        let weight = spill_weight_from_policy(policy, is_hot, operand.kind() != OperandKind::Use);
         u.weight = u16::try_from(weight).expect("weight too large for u16 field");
 
         log::debug!(
@@ -1960,36 +1955,6 @@ impl<'a, F: Function> Env<'a, F> {
         Ok(())
     }
 
-    fn compute_hot_code(&mut self) {
-        // Initialize hot_code to contain inner loops only.
-        let mut header = Block::invalid();
-        let mut backedge = Block::invalid();
-        for block in 0..self.func.blocks() {
-            let block = Block::new(block);
-            let max_backedge = self
-                .func
-                .block_preds(block)
-                .iter()
-                .filter(|b| b.index() >= block.index())
-                .max();
-            if let Some(&b) = max_backedge {
-                header = block;
-                backedge = b;
-            }
-            if block == backedge {
-                // We've traversed a loop body without finding a deeper loop. Mark the whole body
-                // as hot.
-                let from = self.cfginfo.block_entry[header.index()];
-                let to = self.cfginfo.block_exit[backedge.index()].next();
-                let range = CodeRange { from, to };
-                let lr = self.create_liverange(range);
-                self.hot_code
-                    .btree
-                    .insert(LiveRangeKey::from_range(&range), lr);
-            }
-        }
-    }
-
     fn create_bundle(&mut self) -> LiveBundleIndex {
         let bundle = self.bundles.len();
         self.bundles.push(LiveBundle {
@@ -4357,7 +4322,6 @@ impl<'a, F: Function> Env<'a, F> {
 
     pub(crate) fn init(&mut self) -> Result<(), RegAllocError> {
         self.create_pregs_and_vregs();
-        self.compute_hot_code();
         self.compute_liveness()?;
         self.merge_vreg_bundles();
         self.queue_bundles();