Handle up to 2^16 - 1 operands per inst, and error properly on more.

This PR addresses an underlying limitation that became a problem for
Cranelift recently, after folding return-value loads into call
pseudoinsts (thus producing single instructions with many operands): we
only supported 255 operands per instruction.

It appears that `Use` was using a `u8` for the "slot index" for an
operand as a result of earlier optimization/packing efforts. However,
the current shape of the struct leaves a padding byte free, so we should
be able to expand to a `u16` with no loss in performance.

Furthermore, previously we weren't catching the overflow, but rather
were silently wrapping around (eek). This properly returns a
`RegAllocError::TooManyOperands` now if an instruction has too many
operands (more than `u16::MAX`).

Author: cfallin

src/ion/data_structures.rs

@@ -13,13 +13,13 @@
 //! Data structures for backtracking allocator.
 
 use super::liveranges::SpillWeight;
+use crate::Vec2;
 use crate::cfg::{CFGInfo, CFGInfoCtx};
 use crate::index::ContainerComparator;
 use crate::indexset::IndexSet;
-use crate::Vec2;
 use crate::{
-    define_index, Allocation, Block, Bump, Edit, Function, FxHashMap, FxHashSet, MachineEnv,
-    Operand, Output, PReg, ProgPoint, RegClass, VReg,
+    Allocation, Block, Bump, Edit, Function, FxHashMap, FxHashSet, MachineEnv, Operand, Output,
+    PReg, ProgPoint, RegClass, VReg, define_index,
 };
 use alloc::collections::BTreeMap;
 use alloc::collections::VecDeque;
@@ -179,13 +179,13 @@ impl LiveRange {
 pub struct Use {
     pub operand: Operand,
     pub pos: ProgPoint,
-    pub slot: u8,
+    pub slot: u16,
     pub weight: u16,
 }
 
 impl Use {
     #[inline(always)]
-    pub fn new(operand: Operand, pos: ProgPoint, slot: u8) -> Self {
+    pub fn new(operand: Operand, pos: ProgPoint, slot: u16) -> Self {
         Self {
             operand,
             pos,
@@ -315,8 +315,8 @@ pub struct PRegData {
 #[derive(Clone, Debug)]
 pub struct MultiFixedRegFixup {
     pub pos: ProgPoint,
-    pub from_slot: u8,
-    pub to_slot: u8,
+    pub from_slot: u16,
+    pub to_slot: u16,
     pub level: FixedRegFixupLevel,
     pub to_preg: PRegIndex,
     pub vreg: VRegIndex,
@@ -713,11 +713,7 @@ impl InsertedMoves {
     ) {
         trace!(
             "insert_move: pos {:?} prio {:?} from_alloc {:?} to_alloc {:?} to_vreg {:?}",
-            pos,
-            prio,
-            from_alloc,
-            to_alloc,
-            to_vreg
+            pos, prio, from_alloc, to_alloc, to_vreg
         );
         if from_alloc == to_alloc {
             trace!(" -> skipping move with same source and  dest");

src/ion/liveranges.rs

@@ -24,7 +24,8 @@ use crate::{
     Allocation, Block, Function, Inst, InstPosition, Operand, OperandConstraint, OperandKind,
     OperandPos, PReg, ProgPoint, RegAllocError, VReg, VecExt,
 };
-use smallvec::{smallvec, SmallVec};
+use core::convert::TryFrom;
+use smallvec::{SmallVec, smallvec};
 
 /// A spill weight computed for a certain Use.
 #[derive(Clone, Copy, Debug)]
@@ -94,6 +95,10 @@ impl core::ops::Add<SpillWeight> for SpillWeight {
     }
 }
 
+fn slot_idx(i: usize) -> Result<u16, RegAllocError> {
+    u16::try_from(i).map_err(|_| RegAllocError::TooManyOperands)
+}
+
 impl<'a, F: Function> Env<'a, F> {
     pub fn create_pregs_and_vregs(&mut self) {
         // Create PRegs from the env.
@@ -221,9 +226,7 @@ impl<'a, F: Function> Env<'a, F> {
 
         trace!(
             "insert use {:?} into lr {:?} with weight {:?}",
-            u,
-            into,
-            weight,
+            u, into, weight,
         );
 
         // N.B.: we do *not* update `requirement` on the range,
@@ -362,7 +365,7 @@ impl<'a, F: Function> Env<'a, F> {
         Ok(())
     }
 
-    pub fn build_liveranges(&mut self) {
+    pub fn build_liveranges(&mut self) -> Result<(), RegAllocError> {
         // Create Uses and Defs referring to VRegs, and place the Uses
         // in LiveRanges.
         //
@@ -456,9 +459,11 @@ impl<'a, F: Function> Env<'a, F> {
                 let mut reused_input = None;
                 for op in self.func.inst_operands(inst) {
                     if let OperandConstraint::Reuse(i) = op.constraint() {
-                        debug_assert!(self.func.inst_operands(inst)[i]
-                            .as_fixed_nonallocatable()
-                            .is_none());
+                        debug_assert!(
+                            self.func.inst_operands(inst)[i]
+                                .as_fixed_nonallocatable()
+                                .is_none()
+                        );
                         reused_input = Some(self.func.inst_operands(inst)[i].vreg());
                         break;
                     }
@@ -527,14 +532,13 @@ impl<'a, F: Function> Env<'a, F> {
                                             "is downward live, and there is also a ",
                                             "def or clobber at this preg"
                                         ),
-                                        operand,
-                                        preg
+                                        operand, preg
                                     );
                                     let pos = ProgPoint::before(inst);
                                     self.multi_fixed_reg_fixups.push(MultiFixedRegFixup {
                                         pos,
-                                        from_slot: i as u8,
-                                        to_slot: i as u8,
+                                        from_slot: slot_idx(i)?,
+                                        to_slot: slot_idx(i)?,
                                         to_preg: PRegIndex::new(preg.index()),
                                         vreg: VRegIndex::new(operand.vreg().vreg()),
                                         level: FixedRegFixupLevel::Initial,
@@ -638,7 +642,10 @@ impl<'a, F: Function> Env<'a, F> {
                                     live.set(operand.vreg().vreg(), true);
                                 }
                                 // Create the use in the LiveRange.
-                                self.insert_use_into_liverange(lr, Use::new(operand, pos, i as u8));
+                                self.insert_use_into_liverange(
+                                    lr,
+                                    Use::new(operand, pos, slot_idx(i)?),
+                                );
                                 // If def (not mod), this reg is now dead,
                                 // scanning backward; make it so.
                                 if operand.kind() == OperandKind::Def {
@@ -681,7 +688,10 @@ impl<'a, F: Function> Env<'a, F> {
 
                                 trace!("Use of {:?} at {:?} -> {:?}", operand, pos, lr,);
 
-                                self.insert_use_into_liverange(lr, Use::new(operand, pos, i as u8));
+                                self.insert_use_into_liverange(
+                                    lr,
+                                    Use::new(operand, pos, slot_idx(i)?),
+                                );
 
                                 // Add to live-set.
                                 live.set(operand.vreg().vreg(), true);
@@ -756,6 +766,8 @@ impl<'a, F: Function> Env<'a, F> {
         self.output.stats.blockparam_outs_count = self.blockparam_outs.len();
         self.ctx.scratch_vreg_ranges = vreg_ranges;
         self.ctx.scratch_operand_rewrites = operand_rewrites;
+
+        Ok(())
     }
 
     pub fn fixup_multi_fixed_vregs(&mut self) {
@@ -837,9 +849,7 @@ impl<'a, F: Function> Env<'a, F> {
                             let preg_idx = PRegIndex::new(preg.index());
                             trace!(
                                 "at pos {:?}, vreg {:?} has fixed constraint to preg {:?}",
-                                u.pos,
-                                vreg_idx,
-                                preg_idx
+                                u.pos, vreg_idx, preg_idx
                             );
 
                             // FixedStack is incompatible if there are any

src/ion/mod.rs

@@ -83,7 +83,7 @@ impl<'a, F: Function> Env<'a, F> {
     pub(crate) fn init(&mut self) -> Result<(), RegAllocError> {
         self.create_pregs_and_vregs();
         self.compute_liveness()?;
-        self.build_liveranges();
+        self.build_liveranges()?;
         self.fixup_multi_fixed_vregs();
         self.merge_vreg_bundles();
         self.queue_bundles();

src/lib.rs

@@ -1577,6 +1577,9 @@ pub enum RegAllocError {
     /// Too many pinned VRegs + Reg-constrained Operands are live at
     /// once, making allocation impossible.
     TooManyLiveRegs,
+    /// Too many operands on a single instruction (beyond limit of
+    /// 2^16 - 1).
+    TooManyOperands,
 }
 
 impl core::fmt::Display for RegAllocError {