diff --git a/crates/cranelift/src/func_environ.rs b/crates/cranelift/src/func_environ.rs index 854cd4eae429..b829988f8a8d 100644 --- a/crates/cranelift/src/func_environ.rs +++ b/crates/cranelift/src/func_environ.rs @@ -1371,6 +1371,17 @@ impl<'module_environment> FuncEnvironment<'module_environment> { builder.ins().stack_store(vmctx, slot, 0); } } + + pub(crate) fn val_ty_needs_stack_map(&self, ty: WasmValType) -> bool { + match ty { + WasmValType::Ref(r) => self.heap_ty_needs_stack_map(r.heap_type), + _ => false, + } + } + + pub(crate) fn heap_ty_needs_stack_map(&self, ty: WasmHeapType) -> bool { + ty.is_vmgcref_type_and_not_i31() && !ty.is_bottom() + } } impl TranslateTrap for FuncEnvironment<'_> { @@ -2552,13 +2563,7 @@ impl<'module_environment> TargetEnvironment for FuncEnvironment<'module_environm fn reference_type(&self, wasm_ty: WasmHeapType) -> (ir::Type, bool) { let ty = crate::reference_type(wasm_ty, self.pointer_type()); - let needs_stack_map = match wasm_ty.top() { - WasmHeapTopType::Extern | WasmHeapTopType::Any | WasmHeapTopType::Exn => true, - WasmHeapTopType::Func => false, - // TODO(#10248) Once continuations can be stored on the GC heap, we - // will need stack maps for continuation objects. - WasmHeapTopType::Cont => false, - }; + let needs_stack_map = self.heap_ty_needs_stack_map(wasm_ty); (ty, needs_stack_map) } diff --git a/crates/cranelift/src/func_environ/gc/enabled/drc.rs b/crates/cranelift/src/func_environ/gc/enabled/drc.rs index cbda4ebb5fd7..037e09d55146 100644 --- a/crates/cranelift/src/func_environ/gc/enabled/drc.rs +++ b/crates/cranelift/src/func_environ/gc/enabled/drc.rs @@ -261,8 +261,7 @@ impl DrcCompiler { new_val: ir::Value, flags: ir::MemFlags, ) -> WasmResult<()> { - let (ref_ty, needs_stack_map) = func_env.reference_type(ty.heap_type); - debug_assert!(needs_stack_map); + let (ref_ty, _) = func_env.reference_type(ty.heap_type); // Special case for references to uninhabited bottom types: see // `translate_write_gc_reference` for details. @@ -585,8 +584,7 @@ impl GcCompiler for DrcCompiler { assert!(ty.is_vmgcref_type()); - let (reference_type, needs_stack_map) = func_env.reference_type(ty.heap_type); - debug_assert!(needs_stack_map); + let (reference_type, _) = func_env.reference_type(ty.heap_type); // Special case for references to uninhabited bottom types: the // reference must either be nullable and we can just eagerly return @@ -737,8 +735,7 @@ impl GcCompiler for DrcCompiler { ) -> WasmResult<()> { assert!(ty.is_vmgcref_type()); - let (ref_ty, needs_stack_map) = func_env.reference_type(ty.heap_type); - debug_assert!(needs_stack_map); + let (ref_ty, _) = func_env.reference_type(ty.heap_type); // Special case for references to uninhabited bottom types: either the // reference is nullable and we can just eagerly store null into `dst` diff --git a/crates/cranelift/src/translate/code_translator.rs b/crates/cranelift/src/translate/code_translator.rs index 494bb275c71c..64fb55e1e726 100644 --- a/crates/cranelift/src/translate/code_translator.rs +++ b/crates/cranelift/src/translate/code_translator.rs @@ -203,31 +203,47 @@ pub fn translate_operator( Operator::Drop => { environ.stacks.pop1(); } - Operator::Select => { - let (mut arg1, mut arg2, cond) = environ.stacks.pop3(); - if builder.func.dfg.value_type(arg1).is_vector() { - arg1 = optionally_bitcast_vector(arg1, I8X16, builder); - } - if builder.func.dfg.value_type(arg2).is_vector() { - arg2 = optionally_bitcast_vector(arg2, I8X16, builder); - } - environ.stacks.push1(builder.ins().select(cond, arg1, arg2)); + Operator::Nop => { + // We do nothing } - Operator::TypedSelect { ty: _ } => { + Operator::Select + | Operator::TypedSelect { // We ignore the explicit type parameter as it is only needed for // validation, which we require to have been performed before // translation. + ty: _, + } => { let (mut arg1, mut arg2, cond) = environ.stacks.pop3(); + if builder.func.dfg.value_type(arg1).is_vector() { arg1 = optionally_bitcast_vector(arg1, I8X16, builder); } if builder.func.dfg.value_type(arg2).is_vector() { arg2 = optionally_bitcast_vector(arg2, I8X16, builder); } - environ.stacks.push1(builder.ins().select(cond, arg1, arg2)); - } - Operator::Nop => { - // We do nothing + + let val = builder.ins().select(cond, arg1, arg2); + + // If either of the input types need inclusion in stack maps, then + // the result will as well. + // + // Note that we don't need to check whether the result's type needs + // inclusion in stack maps (that would be a conservative over + // approximation) because the input types give us more-precise + // information than the result type does. For example, the result + // does not need inclusion in stack maps in the scenario where both + // inputs are `i31ref`s and the result is an `anyref`. Even though + // `anyref`s normally do require inclusion in stack maps, in this + // particular case, we know that we are dealing with an `anyref` + // that doesn't actually require inclusion. + if operand_types + .iter() + .any(|ty| environ.val_ty_needs_stack_map(*ty)) + { + builder.declare_value_needs_stack_map(val); + } + + environ.stacks.push1(val); } Operator::Unreachable => { environ.trap(builder, crate::TRAP_UNREACHABLE); diff --git a/tests/all/gc.rs b/tests/all/gc.rs index 55b6e9e8b5b6..a49ada513380 100644 --- a/tests/all/gc.rs +++ b/tests/all/gc.rs @@ -1800,3 +1800,80 @@ fn array_init_elem_oom() -> Result<()> { Ok(()) } + +// The result of a `select` instruction with a GC reference type should be +// declared as needing a stack map. +#[test] +#[cfg_attr(miri, ignore)] +fn select_gc_ref_stack_map() -> Result<()> { + let _ = env_logger::try_init(); + + let mut config = Config::new(); + config.wasm_function_references(true); + config.wasm_gc(true); + // Use pooling allocator with a tiny GC heap so that GC is triggered + // frequently and freed memory is reused quickly. + let mut pool = crate::small_pool_config(); + pool.max_memory_size(1 << 16); // 64 KiB + config.allocation_strategy(pool); + + let engine = Engine::new(&config)?; + let module = Module::new( + &engine, + r#" + (module + (type $pair (struct (field (mut i32)))) + (type $arr (array (mut i8))) + + ;; Allocate many objects to fill the GC heap and trigger + ;; collection. After GC frees everything, subsequent + ;; allocations reuse the freed memory. + (func $force_gc + (local i32) + (local.set 0 (i32.const 64)) + (loop $l + (drop (array.new $arr (i32.const 0) (i32.const 1024))) + (local.set 0 (i32.sub (local.get 0) (i32.const 1))) + (br_if $l (local.get 0)) + ) + ) + + (func (export "test") (param $cond i32) (result i32) + ;; The select result stays on the Wasm operand stack (never + ;; stored in a local variable). If the new SSA value created + ;; by select is not declared for stack maps, the GC will + ;; free it. + (select (result (ref null $pair)) + (struct.new $pair (i32.const 111)) + (ref.null $pair) + (local.get $cond) + ) + + ;; This call is a safepoint. The select result is live + ;; across it. The called function triggers GC which will + ;; free the struct if it is incorrectly omitted from the + ;; safepoint's stack map, and the new allocations would + ;; overwrite the freed memory. + (call $force_gc) + + ;; Use the select result. If it was incorrectly freed, then + ;; this will have the wrong value. + (struct.get $pair 0) + ) + ) + "#, + )?; + + let mut store = Store::new(&engine, ()); + let instance = Instance::new(&mut store, &module, &[])?; + let test = instance.get_typed_func::<(i32,), i32>(&mut store, "test")?; + + // Run multiple times to increase chance of triggering GC at the right + // moment. + for _ in 0..30 { + let result = test.call(&mut store, (1,))?; + assert_eq!(result, 111); + } + + Ok(()) +} diff --git a/tests/disas/gc/typed-select-and-stack-maps.wat b/tests/disas/gc/typed-select-and-stack-maps.wat new file mode 100644 index 000000000000..9eaf0158bf47 --- /dev/null +++ b/tests/disas/gc/typed-select-and-stack-maps.wat @@ -0,0 +1,90 @@ +;;! target = "x86_64" +;;! flags = "-W function-references,gc" +;;! test = "optimize" + +(module + (import "" "observe" (func $observe (param anyref))) + (import "" "safepoint" (func $safepoint)) + + (func (param structref i31ref i32) + ;; Select from two types, one of which requires inclusion in stack maps, + ;; resulting in a value that also requires inclusion in stack maps. + (select (result anyref) + (local.get 0) + (local.get 1) + (local.get 2)) + + ;; Make a call, which is a safepoint and has stack maps. + (call $safepoint) + + ;; Observe the result of the select to keep it alive across the call. + call $observe + ) + + (func (param i31ref i31ref i32) + ;; Select from two types that do not require inclusion in stack maps, + ;; resulting in one that (normally) does. In this case, however, we + ;; shouldn't include the value in a stack map, because we know that the + ;; anyref cannot be an instance of a subtype that actually does require + ;; inclusion in stack maps. + (select (result anyref) + (local.get 0) + (local.get 1) + (local.get 2)) + + ;; Make a call, which is a safepoint and has stack maps. + (call $safepoint) + + ;; Observe the result of the select to keep it alive across the call. + call $observe + ) +) +;; function u0:0(i64 vmctx, i64, i32, i32, i32) tail { +;; ss0 = explicit_slot 4, align = 4 +;; gv0 = vmctx +;; gv1 = load.i64 notrap aligned readonly gv0+8 +;; gv2 = load.i64 notrap aligned gv1+24 +;; gv3 = vmctx +;; sig0 = (i64 vmctx, i64) tail +;; sig1 = (i64 vmctx, i64, i32) tail +;; stack_limit = gv2 +;; +;; block0(v0: i64, v1: i64, v2: i32, v3: i32, v4: i32): +;; @0049 v5 = select v4, v2, v3 +;; v14 = stack_addr.i64 ss0 +;; store notrap v5, v14 +;; @004c v8 = load.i64 notrap aligned readonly can_move v0+72 +;; @004c v7 = load.i64 notrap aligned readonly can_move v0+88 +;; @004c call_indirect sig0, v8(v7, v0), stack_map=[i32 @ ss0+0] +;; v12 = load.i32 notrap v14 +;; @004e v11 = load.i64 notrap aligned readonly can_move v0+48 +;; @004e v10 = load.i64 notrap aligned readonly can_move v0+64 +;; @004e call_indirect sig1, v11(v10, v0, v12) +;; @0050 jump block1 +;; +;; block1: +;; @0050 return +;; } +;; +;; function u0:1(i64 vmctx, i64, i32, i32, i32) tail { +;; gv0 = vmctx +;; gv1 = load.i64 notrap aligned readonly gv0+8 +;; gv2 = load.i64 notrap aligned gv1+24 +;; gv3 = vmctx +;; sig0 = (i64 vmctx, i64) tail +;; sig1 = (i64 vmctx, i64, i32) tail +;; stack_limit = gv2 +;; +;; block0(v0: i64, v1: i64, v2: i32, v3: i32, v4: i32): +;; @005c v8 = load.i64 notrap aligned readonly can_move v0+72 +;; @005c v7 = load.i64 notrap aligned readonly can_move v0+88 +;; @005c call_indirect sig0, v8(v7, v0) +;; @005e v11 = load.i64 notrap aligned readonly can_move v0+48 +;; @005e v10 = load.i64 notrap aligned readonly can_move v0+64 +;; @0059 v5 = select v4, v2, v3 +;; @005e call_indirect sig1, v11(v10, v0, v5) +;; @0060 jump block1 +;; +;; block1: +;; @0060 return +;; }