Add limits to allocatable registers

src/checker.rs

@@ -104,6 +104,7 @@ use alloc::vec::Vec;
 use alloc::{format, vec};
 use core::default::Default;
 use core::hash::Hash;
+use core::ops::Range;
 use core::result::Result;
 use smallvec::{smallvec, SmallVec};
 
@@ -166,6 +167,12 @@ pub enum CheckerError {
         into: Allocation,
         from: Allocation,
     },
+    AllocationOutsideLimit {
+        inst: Inst,
+        op: Operand,
+        alloc: Allocation,
+        range: Range<usize>,
+    },
 }
 
 /// Abstract state for an allocation.
@@ -669,6 +676,18 @@ impl CheckerState {
                     });
                 }
             }
+            OperandConstraint::Limit(max) => {
+                if let Some(preg) = alloc.as_reg() {
+                    if preg.hw_enc() >= max {
+                        return Err(CheckerError::AllocationOutsideLimit {
+                            inst,
+                            op,
+                            alloc,
+                            range: (0..max),
+                        });
+                    }
+                }
+            }
         }
         Ok(())
     }

src/fastalloc/mod.rs

@@ -509,6 +509,9 @@ impl<'a, F: Function> Env<'a, F> {
             }
 
             OperandConstraint::Stack => self.edits.is_stack(alloc),
+            OperandConstraint::Limit(_) => {
+                todo!("limit constraints are not yet supported in fastalloc")
+            }
         }
     }
 
@@ -667,6 +670,9 @@ impl<'a, F: Function> Env<'a, F> {
             }
 
             OperandConstraint::Stack => Allocation::stack(self.get_spillslot(op.vreg())),
+            OperandConstraint::Limit(_) => {
+                todo!("limit constraints are not yet supported in fastalloc")
+            }
         };
         self.allocs[(inst.index(), op_idx)] = new_alloc;
         Ok(new_alloc)

src/ion/data_structures.rs

@@ -203,6 +203,7 @@ pub struct LiveBundle {
     pub allocation: Allocation,
     pub prio: u32, // recomputed after every bulk update
     pub spill_weight_and_props: u32,
+    pub limit: Option<u8>,
 }
 
 pub const BUNDLE_MAX_SPILL_WEIGHT: u32 = (1 << 28) - 1;
@@ -413,6 +414,7 @@ impl LiveBundles {
             spillset: SpillSetIndex::invalid(),
             prio: 0,
             spill_weight_and_props: 0,
+            limit: None,
         })
     }
 }

src/ion/liveranges.rs

@@ -25,6 +25,7 @@ use crate::{
     OperandPos, PReg, ProgPoint, RegAllocError, VReg, VecExt,
 };
 use core::convert::TryFrom;
+use core::usize;
 use smallvec::{smallvec, SmallVec};
 
 /// A spill weight computed for a certain Use.
@@ -804,6 +805,8 @@ impl<'a, F: Function> Env<'a, F> {
                     let mut num_fixed_stack = 0;
                     let mut first_reg_slot = None;
                     let mut first_stack_slot = None;
+                    let mut min_limit = usize::MAX;
+                    let mut max_fixed_reg = usize::MIN;
                     for u in uses.iter() {
                         match u.operand.constraint() {
                             OperandConstraint::Any => {
@@ -814,7 +817,13 @@ impl<'a, F: Function> Env<'a, F> {
                                 first_reg_slot.get_or_insert(u.slot);
                                 requires_reg = true;
                             }
+                            OperandConstraint::Limit(max) => {
+                                first_reg_slot.get_or_insert(u.slot);
+                                min_limit = min_limit.min(max);
+                                requires_reg = true;
+                            }
                             OperandConstraint::FixedReg(preg) => {
+                                max_fixed_reg = max_fixed_reg.max(preg.hw_enc());
                                 if self.ctx.pregs[preg.index()].is_stack {
                                     num_fixed_stack += 1;
                                     first_stack_slot.get_or_insert(u.slot);
@@ -834,6 +843,7 @@ impl<'a, F: Function> Env<'a, F> {
                     // Fast path if there are no conflicts.
                     if num_fixed_reg + num_fixed_stack <= 1
                         && !(requires_reg && num_fixed_stack != 0)
+                        && max_fixed_reg < min_limit
                     {
                         continue;
                     }
@@ -867,6 +877,7 @@ impl<'a, F: Function> Env<'a, F> {
                             // skip this edit.
                             if !(requires_reg && self.ctx.pregs[preg.index()].is_stack)
                                 && *first_preg.get_or_insert(preg) == preg
+                                && preg.hw_enc() < min_limit
                             {
                                 continue;
                             }

src/ion/merge.rs

@@ -12,17 +12,33 @@
 
 //! Bundle merging.
 
-use super::{Env, LiveBundleIndex, SpillSet, SpillSlotIndex, VRegIndex};
-use crate::{
-    ion::{
-        data_structures::{BlockparamOut, CodeRange},
-        LiveRangeList,
-    },
-    Function, Inst, OperandConstraint, OperandKind, PReg, ProgPoint,
+use crate::ion::data_structures::{
+    BlockparamOut, CodeRange, Env, LiveBundleIndex, LiveRangeList, SpillSet, SpillSlotIndex,
+    VRegIndex,
 };
+use crate::{Function, Inst, OperandConstraint, OperandKind, PReg, ProgPoint};
 use alloc::format;
+use core::convert::TryFrom;
 
 impl<'a, F: Function> Env<'a, F> {
+    fn merge_bundle_properties(&mut self, from: LiveBundleIndex, to: LiveBundleIndex) {
+        if self.bundles[from].cached_fixed() {
+            self.bundles[to].set_cached_fixed();
+        }
+        if self.bundles[from].cached_fixed_def() {
+            self.bundles[to].set_cached_fixed_def();
+        }
+        if self.bundles[from].cached_stack() {
+            self.bundles[to].set_cached_stack();
+        }
+        if let Some(theirs) = self.bundles[from].limit {
+            match self.bundles[to].limit {
+                Some(ours) => self.bundles[to].limit = Some(ours.min(theirs)),
+                None => self.bundles[to].limit = Some(theirs),
+            }
+        }
+    }
+
     pub fn merge_bundles(&mut self, from: LiveBundleIndex, to: LiveBundleIndex) -> bool {
         if from == to {
             // Merge bundle into self -- trivial merge.
@@ -114,8 +130,10 @@ impl<'a, F: Function> Env<'a, F> {
         // Check for a requirements conflict.
         if self.bundles[from].cached_stack()
             || self.bundles[from].cached_fixed()
+            || self.bundles[from].limit.is_some()
             || self.bundles[to].cached_stack()
             || self.bundles[to].cached_fixed()
+            || self.bundles[to].limit.is_some()
         {
             if self.merge_bundle_requirements(from, to).is_err() {
                 trace!(" -> conflicting requirements; aborting merge");
@@ -157,17 +175,7 @@ impl<'a, F: Function> Env<'a, F> {
                 }
             }
             self.bundles[to].ranges = list;
-
-            if self.bundles[from].cached_stack() {
-                self.bundles[to].set_cached_stack();
-            }
-            if self.bundles[from].cached_fixed() {
-                self.bundles[to].set_cached_fixed();
-            }
-            if self.bundles[from].cached_fixed_def() {
-                self.bundles[to].set_cached_fixed_def();
-            }
-
+            self.merge_bundle_properties(from, to);
             return true;
         }
 
@@ -243,15 +251,7 @@ impl<'a, F: Function> Env<'a, F> {
             *to_range = to_range.join(from_range);
         }
 
-        if self.bundles[from].cached_stack() {
-            self.bundles[to].set_cached_stack();
-        }
-        if self.bundles[from].cached_fixed() {
-            self.bundles[to].set_cached_fixed();
-        }
-        if self.bundles[from].cached_fixed_def() {
-            self.bundles[to].set_cached_fixed_def();
-        }
+        self.merge_bundle_properties(from, to);
 
         true
     }
@@ -283,16 +283,29 @@ impl<'a, F: Function> Env<'a, F> {
             let mut fixed = false;
             let mut fixed_def = false;
             let mut stack = false;
+            let mut limit: Option<u8> = None;
             for entry in &self.bundles[bundle].ranges {
                 for u in &self.ranges[entry.index].uses {
-                    if let OperandConstraint::FixedReg(_) = u.operand.constraint() {
-                        fixed = true;
-                        if u.operand.kind() == OperandKind::Def {
-                            fixed_def = true;
+                    use OperandConstraint::*;
+                    match u.operand.constraint() {
+                        FixedReg(_) => {
+                            fixed = true;
+                            if u.operand.kind() == OperandKind::Def {
+                                fixed_def = true;
+                            }
+                        }
+                        Stack => stack = true,
+                        Limit(current) => {
+                            let current = u8::try_from(current)
+                                .expect("the current limit is too large to fit in a u8");
+                            match limit {
+                                Some(prev) => limit = Some(prev.min(current)),
+                                None => limit = Some(current),
+                            }
+                        }
+                        Any | Reg | Reuse(_) => {
+                            continue;
                         }
-                    }
-                    if let OperandConstraint::Stack = u.operand.constraint() {
-                        stack = true;
                     }
                     if fixed && stack && fixed_def {
                         break;
@@ -308,6 +321,7 @@ impl<'a, F: Function> Env<'a, F> {
             if stack {
                 self.bundles[bundle].set_cached_stack();
             }
+            self.bundles[bundle].limit = limit;
 
             // Create a spillslot for this bundle.
             let reg = self.vreg(vreg);
@@ -383,6 +397,32 @@ impl<'a, F: Function> Env<'a, F> {
         total
     }
 
+    pub fn compute_bundle_limit(&self, bundle: LiveBundleIndex) -> Option<u8> {
+        let mut limit: Option<u8> = None;
+        for entry in &self.bundles[bundle].ranges {
+            for u in &self.ranges[entry.index].uses {
+                use OperandConstraint::*;
+                match u.operand.constraint() {
+                    Limit(current) => {
+                        let current = u8::try_from(current)
+                            .expect("the current limit is too large to fit in a u8");
+                        match limit {
+                            Some(prev) => limit = Some(prev.min(current)),
+                            None => limit = Some(current),
+                        }
+                    }
+                    FixedReg(_) | Stack => {
+                        break;
+                    }
+                    Any | Reg | Reuse(_) => {
+                        continue;
+                    }
+                }
+            }
+        }
+        limit
+    }
+
     pub fn queue_bundles(&mut self) {
         for bundle in 0..self.bundles.len() {
             trace!("enqueueing bundle{}", bundle);

src/ion/moves.rs

@@ -892,7 +892,10 @@ impl<'a, F: Function> Env<'a, F> {
                 }
 
                 let resolved = parallel_moves.resolve();
-                let mut scratch_iter = RegTraversalIter::new(self.env, regclass, None, None, 0);
+                let mut scratch_iter = RegTraversalIter::new(
+                    self.env, regclass, None, None, 0,
+                    None, // We assume there is no limit on the set of registers available for moves.
+                );
                 let mut dedicated_scratch = self.env.scratch_by_class[regclass as usize];
                 let key = LiveRangeKey::from_range(&CodeRange {
                     from: pos_prio.pos,

src/ion/process.rs

@@ -274,6 +274,7 @@ impl<'a, F: Function> Env<'a, F> {
         let first_range_data = &self.ctx.ranges[first_range];
 
         self.ctx.bundles[bundle].prio = self.compute_bundle_prio(bundle);
+        self.ctx.bundles[bundle].limit = self.compute_bundle_limit(bundle);
 
         if first_range_data.vreg.is_invalid() {
             trace!("  -> no vreg; minimal and fixed");
@@ -1036,7 +1037,7 @@ impl<'a, F: Function> Env<'a, F> {
 
             let fixed_preg = match req {
                 Requirement::FixedReg(preg) | Requirement::FixedStack(preg) => Some(preg),
-                Requirement::Register => None,
+                Requirement::Register | Requirement::Limit(..) => None,
                 Requirement::Stack => {
                     // If we must be on the stack, mark our spillset
                     // as required immediately.
@@ -1072,9 +1073,15 @@ impl<'a, F: Function> Env<'a, F> {
                 + bundle.index();
 
             self.ctx.output.stats.process_bundle_reg_probe_start_any += 1;
-            for preg in
-                RegTraversalIter::new(self.env, class, fixed_preg, hint.as_valid(), scan_offset)
-            {
+            let limit = self.bundles[bundle].limit.map(|l| l as usize);
+            for preg in RegTraversalIter::new(
+                self.env,
+                class,
+                fixed_preg,
+                hint.as_valid(),
+                scan_offset,
+                limit,
+            ) {
                 self.ctx.output.stats.process_bundle_reg_probes_any += 1;
                 let preg_idx = PRegIndex::new(preg.index());
                 trace!("trying preg {:?}", preg_idx);
@@ -1200,7 +1207,7 @@ impl<'a, F: Function> Env<'a, F> {
                     || lowest_cost_evict_conflict_cost.is_none()
                     || lowest_cost_evict_conflict_cost.unwrap() >= our_spill_weight)
             {
-                if let Requirement::Register = req {
+                if matches!(req, Requirement::Register | Requirement::Limit(_)) {
                     // Check if this is a too-many-live-registers situation.
                     let range = self.ctx.bundles[bundle].ranges[0].range;
                     trace!("checking for too many live regs");
@@ -1240,6 +1247,15 @@ impl<'a, F: Function> Env<'a, F> {
                                 fixed_assigned += 1;
                             }
                         }
+
+                        // We also need to discard any registers that do not fit
+                        // under the limit--we cannot allocate to them.
+                        if let Requirement::Limit(limit) = req {
+                            if preg.hw_enc() >= limit as usize {
+                                continue;
+                            }
+                        }
+
                         total_regs += 1;
                     }
                     trace!(