Cranelift: return programmatic error rather than panic when temporari…

…es run out. (bytecodealliance#7114)

* Cranelift: return programmatic error rather than panic when temporaries run out.

Cranelift currently has a limit of `2^21` vregs per function body in
VCode. This is a consequence of (performance-motivated) bitpacking of
`Operand`s into `u32`s in regalloc2.

As a result of this, it is possible to produce a function body that will
fail to compile by running out of vreg temporaries during lowering.
Currently, this results in a panic. Ideally, we would propagate the
error upward and return it programmatically.

This PR does that, with a "deferred error" mechanism. A cleaner solution
would be to properly thread the `Result` types through all layers of
lowering. However, that would require supporting `Result`s in ISLE, and
that is a deeper language-design and `islec`-hacking question that I
think we can tackle later if we continue to see motivating cases.

The deferral works by returning a valid but bogus `ValueReg`s to the
lowering rule, but storing the error and checking for it in the toplevel
lowering loop. (Note that we have to return a bogus `v0` rather than
`VReg::invalid()`, because the latter causes the `ValueRegs` to think
there is no register provided.)

This PR also includes a test at the Wasmtime level. Note that it takes
~22s to run on my (relatively fast) laptop, because it has to run until
it runs out of VRegs in a debug build of the compiler. We could remove
the test if we feel we're otherwise confident.

Thanks to Venkkatesh Sekar for reporting this issue! The integration
test uses one of the example inputs from the report.

* Review feedback.

* Handle alloc after a deferred error occurs.

* Add uncommitted fix for previous.

cranelift/codegen/src/machinst/abi.rs

@@ -1559,7 +1559,7 @@ impl<M: ABIMachineSpec> Callee<M> {
                 // We need to dereference the pointer.
                 let base = match &pointer {
                     &ABIArgSlot::Reg { reg, ty, .. } => {
-                        let tmp = vregs.alloc(ty).unwrap().only_reg().unwrap();
+                        let tmp = vregs.alloc_with_deferred_error(ty).only_reg().unwrap();
                         self.reg_args.push(ArgPair {
                             vreg: Writable::from_reg(tmp),
                             preg: reg.into(),
@@ -1621,9 +1621,10 @@ impl<M: ABIMachineSpec> Callee<M> {
                                     if n < word_bits =>
                                 {
                                     let signed = ext == ir::ArgumentExtension::Sext;
-                                    let dst = writable_value_regs(vregs.alloc(ty).unwrap())
-                                        .only_reg()
-                                        .unwrap();
+                                    let dst =
+                                        writable_value_regs(vregs.alloc_with_deferred_error(ty))
+                                            .only_reg()
+                                            .unwrap();
                                     ret.push(M::gen_extend(
                                         dst, from_reg, signed, from_bits,
                                         /* to_bits = */ word_bits,
@@ -1664,9 +1665,10 @@ impl<M: ABIMachineSpec> Callee<M> {
                                 {
                                     assert_eq!(M::word_reg_class(), from_reg.class());
                                     let signed = ext == ir::ArgumentExtension::Sext;
-                                    let dst = writable_value_regs(vregs.alloc(ty).unwrap())
-                                        .only_reg()
-                                        .unwrap();
+                                    let dst =
+                                        writable_value_regs(vregs.alloc_with_deferred_error(ty))
+                                            .only_reg()
+                                            .unwrap();
                                     ret.push(M::gen_extend(
                                         dst, from_reg, signed, from_bits,
                                         /* to_bits = */ word_bits,

cranelift/codegen/src/machinst/lower.rs

@@ -1069,6 +1069,12 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
             }
 
             self.finish_bb();
+
+            // Check for any deferred vreg-temp allocation errors, and
+            // bubble one up at this time if it exists.
+            if let Some(e) = self.vregs.take_deferred_error() {
+                return Err(e);
+            }
         }
 
         // Now that we've emitted all instructions into the
@@ -1319,7 +1325,7 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
 impl<'func, I: VCodeInst> Lower<'func, I> {
     /// Get a new temp.
     pub fn alloc_tmp(&mut self, ty: Type) -> ValueRegs<Writable<Reg>> {
-        writable_value_regs(self.vregs.alloc(ty).unwrap())
+        writable_value_regs(self.vregs.alloc_with_deferred_error(ty))
     }
 
     /// Emit a machine instruction.

cranelift/codegen/src/machinst/vcode.rs

@@ -1456,6 +1456,12 @@ pub struct VRegAllocator<I> {
     /// reftype-status of each vreg) for efficient iteration.
     reftyped_vregs: Vec<VReg>,
 
+    /// A deferred error, to be bubbled up to the top level of the
+    /// lowering algorithm. We take this approach because we cannot
+    /// currently propagate a `Result` upward through ISLE code (the
+    /// lowering rules) or some ABI code.
+    deferred_error: Option<CodegenError>,
+
     /// The type of instruction that this allocator makes registers for.
     _inst: core::marker::PhantomData<I>,
 }
@@ -1468,12 +1474,16 @@ impl<I: VCodeInst> VRegAllocator<I> {
             vreg_types: vec![],
             reftyped_vregs_set: FxHashSet::default(),
             reftyped_vregs: vec![],
+            deferred_error: None,
             _inst: core::marker::PhantomData::default(),
         }
     }
 
     /// Allocate a fresh ValueRegs.
     pub fn alloc(&mut self, ty: Type) -> CodegenResult<ValueRegs<Reg>> {
+        if self.deferred_error.is_some() {
+            return Err(CodegenError::CodeTooLarge);
+        }
         let v = self.next_vreg;
         let (regclasses, tys) = I::rc_for_type(ty)?;
         self.next_vreg += regclasses.len();
@@ -1495,6 +1505,38 @@ impl<I: VCodeInst> VRegAllocator<I> {
         Ok(regs)
     }
 
+    /// Allocate a fresh ValueRegs, deferring any out-of-vregs
+    /// errors. This is useful in places where we cannot bubble a
+    /// `CodegenResult` upward easily, and which are known to be
+    /// invoked from within the lowering loop that checks the deferred
+    /// error status below.
+    pub fn alloc_with_deferred_error(&mut self, ty: Type) -> ValueRegs<Reg> {
+        match self.alloc(ty) {
+            Ok(x) => x,
+            Err(e) => {
+                self.deferred_error = Some(e);
+                self.bogus_for_deferred_error(ty)
+            }
+        }
+    }
+
+    /// Take any deferred error that was accumulated by `alloc_with_deferred_error`.
+    pub fn take_deferred_error(&mut self) -> Option<CodegenError> {
+        self.deferred_error.take()
+    }
+
+    /// Produce an bogus VReg placeholder with the proper number of
+    /// registers for the given type. This is meant to be used with
+    /// deferred allocation errors (see `Lower::alloc_tmp()`).
+    fn bogus_for_deferred_error(&self, ty: Type) -> ValueRegs<Reg> {
+        let (regclasses, _tys) = I::rc_for_type(ty).expect("must have valid type");
+        match regclasses {
+            &[rc0] => ValueRegs::one(VReg::new(0, rc0).into()),
+            &[rc0, rc1] => ValueRegs::two(VReg::new(0, rc0).into(), VReg::new(1, rc1).into()),
+            _ => panic!("Value must reside in 1 or 2 registers"),
+        }
+    }
+
     /// Set the type of this virtual register.
     pub fn set_vreg_type(&mut self, vreg: VirtualReg, ty: Type) {
         if self.vreg_types.len() <= vreg.index() {

tests/all/code_too_large.rs

@@ -0,0 +1,36 @@
+#![cfg(not(miri))]
+
+use anyhow::Result;
+use wasmtime::*;
+
+#[test]
+fn code_too_large_without_panic() -> Result<()> {
+    const N: usize = 120000;
+
+    // Build a module with a function whose body will allocate too many
+    // temporaries for our current (Cranelift-based) compiler backend to
+    // handle. This test ensures that we propagate the failure upward
+    // and return it programmatically, rather than panic'ing. If we ever
+    // improve our compiler backend to actually handle such a large
+    // function body, we'll need to increase the limits here too!
+    let mut s = String::new();
+    s.push_str("(module\n");
+    s.push_str("(table 1 1 funcref)\n");
+    s.push_str("(func (export \"\") (result i32)\n");
+    s.push_str("i32.const 0\n");
+    for _ in 0..N {
+        s.push_str("table.get 0\n");
+        s.push_str("ref.is_null\n");
+    }
+    s.push_str("))\n");
+
+    let store = Store::<()>::default();
+    let result = Module::new(store.engine(), &s);
+    match result {
+        Err(e) => assert!(e
+            .to_string()
+            .starts_with("Compilation error: Code for function is too large")),
+        Ok(_) => panic!("Please adjust limits to make the module too large to compile!"),
+    }
+    Ok(())
+}

tests/all/main.rs

@@ -3,6 +3,7 @@
 mod async_functions;
 mod call_hook;
 mod cli_tests;
+mod code_too_large;
 mod component_model;
 mod coredump;
 mod custom_signal_handler;