From fea4ccb690c9c7dd36d109764cbe45eb9f0be1da Mon Sep 17 00:00:00 2001 From: Chris Fallin Date: Fri, 29 Sep 2023 11:19:03 -0700 Subject: [PATCH 1/4] 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. --- cranelift/codegen/src/machinst/lower.rs | 24 ++++++++++++++++- cranelift/codegen/src/machinst/vcode.rs | 12 +++++++++ tests/all/code_too_large.rs | 36 +++++++++++++++++++++++++ tests/all/main.rs | 1 + 4 files changed, 72 insertions(+), 1 deletion(-) create mode 100644 tests/all/code_too_large.rs diff --git a/cranelift/codegen/src/machinst/lower.rs b/cranelift/codegen/src/machinst/lower.rs index fe200e5d66f0..4ea6667baa8f 100644 --- a/cranelift/codegen/src/machinst/lower.rs +++ b/cranelift/codegen/src/machinst/lower.rs @@ -18,6 +18,7 @@ use crate::machinst::{ SigSet, VCode, VCodeBuilder, VCodeConstant, VCodeConstantData, VCodeConstants, VCodeInst, ValueRegs, Writable, }; +use crate::CodegenError; use crate::{trace, CodegenResult}; use alloc::vec::Vec; use cranelift_control::ControlPlane; @@ -163,6 +164,13 @@ pub struct Lower<'func, I: VCodeInst> { /// VReg allocation context, given to the vcode field at build time to finalize the vcode. vregs: VRegAllocator, + /// A deferred error from vreg allocation, 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), and temp allocations + /// typically happen as a result of particular lowering rules. + vreg_alloc_error: Option, + /// Mapping from `Value` (SSA value in IR) to virtual register. value_regs: SecondaryMap>, @@ -437,6 +445,7 @@ impl<'func, I: VCodeInst> Lower<'func, I> { allocatable: PRegSet::from(machine_env), vcode, vregs, + vreg_alloc_error: None, value_regs, sret_reg, block_end_colors, @@ -1077,6 +1086,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.vreg_alloc_error { + return Err(e); + } + // Now that we've emitted all instructions into the // VCodeBuilder, let's build the VCode. let vcode = self.vcode.build(self.allocatable, self.vregs); @@ -1325,7 +1340,14 @@ 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_value_regs(self.vregs.alloc(ty).unwrap()) + let allocated = match self.vregs.alloc(ty) { + Ok(x) => x, + Err(e) => { + self.vreg_alloc_error = Some(e); + self.vregs.invalid(ty) + } + }; + writable_value_regs(allocated) } /// Emit a machine instruction. diff --git a/cranelift/codegen/src/machinst/vcode.rs b/cranelift/codegen/src/machinst/vcode.rs index b0a8d58af14a..b1da1145ce05 100644 --- a/cranelift/codegen/src/machinst/vcode.rs +++ b/cranelift/codegen/src/machinst/vcode.rs @@ -1493,6 +1493,18 @@ impl VRegAllocator { Ok(regs) } + /// 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()`). + pub fn invalid(&self, ty: Type) -> ValueRegs { + 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() { diff --git a/tests/all/code_too_large.rs b/tests/all/code_too_large.rs new file mode 100644 index 000000000000..c7c5080dcbce --- /dev/null +++ b/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(()) +} diff --git a/tests/all/main.rs b/tests/all/main.rs index 1c3e6d4af483..54c8229b8eb1 100644 --- a/tests/all/main.rs +++ b/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; From 776d6b8737ad8071b8cbb87bfa6b1dc4ca6f865e Mon Sep 17 00:00:00 2001 From: Chris Fallin Date: Fri, 29 Sep 2023 12:40:26 -0700 Subject: [PATCH 2/4] Review feedback. --- cranelift/codegen/src/machinst/abi.rs | 16 ++++++++------ cranelift/codegen/src/machinst/lower.rs | 28 +++++------------------- cranelift/codegen/src/machinst/vcode.rs | 29 ++++++++++++++++++++++++- 3 files changed, 43 insertions(+), 30 deletions(-) diff --git a/cranelift/codegen/src/machinst/abi.rs b/cranelift/codegen/src/machinst/abi.rs index 950096511d9b..c979e1f52567 100644 --- a/cranelift/codegen/src/machinst/abi.rs +++ b/cranelift/codegen/src/machinst/abi.rs @@ -1545,7 +1545,7 @@ impl Callee { // 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(), @@ -1607,9 +1607,10 @@ impl Callee { 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, @@ -1650,9 +1651,10 @@ impl Callee { { 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, diff --git a/cranelift/codegen/src/machinst/lower.rs b/cranelift/codegen/src/machinst/lower.rs index 4ea6667baa8f..b208c48eeb18 100644 --- a/cranelift/codegen/src/machinst/lower.rs +++ b/cranelift/codegen/src/machinst/lower.rs @@ -18,7 +18,6 @@ use crate::machinst::{ SigSet, VCode, VCodeBuilder, VCodeConstant, VCodeConstantData, VCodeConstants, VCodeInst, ValueRegs, Writable, }; -use crate::CodegenError; use crate::{trace, CodegenResult}; use alloc::vec::Vec; use cranelift_control::ControlPlane; @@ -164,13 +163,6 @@ pub struct Lower<'func, I: VCodeInst> { /// VReg allocation context, given to the vcode field at build time to finalize the vcode. vregs: VRegAllocator, - /// A deferred error from vreg allocation, 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), and temp allocations - /// typically happen as a result of particular lowering rules. - vreg_alloc_error: Option, - /// Mapping from `Value` (SSA value in IR) to virtual register. value_regs: SecondaryMap>, @@ -445,7 +437,6 @@ impl<'func, I: VCodeInst> Lower<'func, I> { allocatable: PRegSet::from(machine_env), vcode, vregs, - vreg_alloc_error: None, value_regs, sret_reg, block_end_colors, @@ -1084,12 +1075,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.vreg_alloc_error { - return Err(e); + // 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 @@ -1340,14 +1331,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> { - let allocated = match self.vregs.alloc(ty) { - Ok(x) => x, - Err(e) => { - self.vreg_alloc_error = Some(e); - self.vregs.invalid(ty) - } - }; - writable_value_regs(allocated) + writable_value_regs(self.vregs.alloc_with_deferred_error(ty)) } /// Emit a machine instruction. diff --git a/cranelift/codegen/src/machinst/vcode.rs b/cranelift/codegen/src/machinst/vcode.rs index b1da1145ce05..f1e21b6b483b 100644 --- a/cranelift/codegen/src/machinst/vcode.rs +++ b/cranelift/codegen/src/machinst/vcode.rs @@ -1454,6 +1454,12 @@ pub struct VRegAllocator { /// reftype-status of each vreg) for efficient iteration. reftyped_vregs: Vec, + /// 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, + /// The type of instruction that this allocator makes registers for. _inst: core::marker::PhantomData, } @@ -1466,6 +1472,7 @@ impl VRegAllocator { vreg_types: vec![], reftyped_vregs_set: FxHashSet::default(), reftyped_vregs: vec![], + deferred_error: None, _inst: core::marker::PhantomData::default(), } } @@ -1493,10 +1500,30 @@ impl VRegAllocator { 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 { + 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 { + 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()`). - pub fn invalid(&self, ty: Type) -> ValueRegs { + fn bogus_for_deferred_error(&self, ty: Type) -> ValueRegs { let (regclasses, _tys) = I::rc_for_type(ty).expect("must have valid type"); match regclasses { &[rc0] => ValueRegs::one(VReg::new(0, rc0).into()), From 1e6d4ff18af1c5cf30292e6a12ac2b5a6f293dc4 Mon Sep 17 00:00:00 2001 From: Chris Fallin Date: Fri, 29 Sep 2023 12:43:58 -0700 Subject: [PATCH 3/4] Handle alloc after a deferred error occurs. --- cranelift/codegen/src/machinst/vcode.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/cranelift/codegen/src/machinst/vcode.rs b/cranelift/codegen/src/machinst/vcode.rs index f1e21b6b483b..88fdcf6d7b71 100644 --- a/cranelift/codegen/src/machinst/vcode.rs +++ b/cranelift/codegen/src/machinst/vcode.rs @@ -1479,6 +1479,9 @@ impl VRegAllocator { /// Allocate a fresh ValueRegs. pub fn alloc(&mut self, ty: Type) -> CodegenResult> { + if self.deferred_error.is_some() { + return Err(self.deferred_error.unwrap()); + } let v = self.next_vreg; let (regclasses, tys) = I::rc_for_type(ty)?; self.next_vreg += regclasses.len(); From 0427e7ab8b70ef85bb1090c2cd9f4299449d7428 Mon Sep 17 00:00:00 2001 From: Chris Fallin Date: Fri, 29 Sep 2023 13:02:48 -0700 Subject: [PATCH 4/4] Add uncommitted fix for previous. --- cranelift/codegen/src/machinst/vcode.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cranelift/codegen/src/machinst/vcode.rs b/cranelift/codegen/src/machinst/vcode.rs index 88fdcf6d7b71..47216859f0e5 100644 --- a/cranelift/codegen/src/machinst/vcode.rs +++ b/cranelift/codegen/src/machinst/vcode.rs @@ -1480,7 +1480,7 @@ impl VRegAllocator { /// Allocate a fresh ValueRegs. pub fn alloc(&mut self, ty: Type) -> CodegenResult> { if self.deferred_error.is_some() { - return Err(self.deferred_error.unwrap()); + return Err(CodegenError::CodeTooLarge); } let v = self.next_vreg; let (regclasses, tys) = I::rc_for_type(ty)?;