Skip to content

Make clearing and restoring the may_leave flag more amenable to opt… #13713

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
fitzgen wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Make clearing and restoring the may_leave flag more amenable to opt… #13713

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions crates/cranelift/src/alias_region.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1631,8 +1631,8 @@ impl AliasRegions<VMComponentOffsets<u8>> {
)
}

/// Load a component instance's flags from the `VMComponentContext`.
pub fn vmcomponent_instance_flags(
/// Load a component instance's `may_leave` flag from the `VMComponentContext`.
pub fn vmcomponent_instance_may_leave(
&mut self,
cursor: &mut FuncCursor<'_>,
vmctx: ir::Value,
Expand All @@ -1643,7 +1643,7 @@ impl AliasRegions<VMComponentOffsets<u8>> {
ir::types::I32,
ir::MemFlagsData::trusted(),
vmctx,
self.offsets.instance_flags(instance),
self.offsets.may_leave(instance),
)
}
}
13 changes: 4 additions & 9 deletions crates/cranelift/src/compiler/component.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1113,9 +1113,8 @@ impl<'a> TrampolineCompiler<'a> {
// run_destructor_block:
// ;; test may_leave, but only if the component instances
// ;; differ
// flags = load.i32 vmctx+$instance_flags_offset
// masked = band flags, $FLAG_MAY_LEAVE
// trapz masked, $TRAP_CANNOT_LEAVE_COMPONENT
// may_leave = load.i32 vmctx+$instance_flags_offset
// trapz may_leave, $TRAP_CANNOT_LEAVE_COMPONENT
//
// ;; set may_block to false, saving the old value to restore
// ;; later, but only if the component instances differ and
Expand Down Expand Up @@ -1527,17 +1526,13 @@ impl<'a> TrampolineCompiler<'a> {
fn check_may_leave_instance(&mut self, instance: RuntimeComponentInstanceIndex) {
let vmctx = self.builder.func.dfg.block_params(self.block0)[0];

let flags = self.alias_regions.vmcomponent_instance_flags(
let may_leave = self.alias_regions.vmcomponent_instance_may_leave(
&mut self.builder.cursor(),
vmctx,
instance,
);
let may_leave_bit = self
.builder
.ins()
.band_imm_u(flags, i64::from(FLAG_MAY_LEAVE));
let (mut traps, builder) = self.traps();
traps.trapz(builder, may_leave_bit, TRAP_CANNOT_LEAVE_COMPONENT);
traps.trapz(builder, may_leave, TRAP_CANNOT_LEAVE_COMPONENT);
}

fn traps(
Expand Down
18 changes: 7 additions & 11 deletions crates/environ/src/component/vmcomponent_offsets.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
// magic: u32,
// builtins: &'static VMComponentBuiltins,
// limits: *const VMStoreContext,
// flags: [VMGlobalDefinition; component.num_runtime_component_instances],
// may_leave: [VMGlobalDefinition; component.num_runtime_component_instances],
// task_may_block: u32,
// trampoline_func_refs: [VMFuncRef; component.num_trampolines],
// unsafe_intrinsics: [VMFuncRef; component.num_unsafe_intrinsics],
Expand All @@ -26,10 +26,6 @@ use crate::component::*;
/// double-checked on `VMComponentContext::from_opaque`.
pub const VMCOMPONENT_MAGIC: u32 = u32::from_le_bytes(*b"comp");

/// Flag for the `VMComponentContext::flags` field which corresponds to the
/// canonical ABI flag `may_leave`
pub const FLAG_MAY_LEAVE: i32 = 1 << 0;

/// Runtime offsets within a `VMComponentContext` for a specific component.
#[derive(Debug, Clone, Copy)]
pub struct VMComponentOffsets<P> {
Expand Down Expand Up @@ -63,7 +59,7 @@ pub struct VMComponentOffsets<P> {
magic: u32,
builtins: u32,
vm_store_context: u32,
flags: u32,
may_leave: u32,
task_may_block: u32,
trampoline_func_refs: u32,
intrinsic_func_refs: u32,
Expand Down Expand Up @@ -132,7 +128,7 @@ impl<P: PtrSize> VMComponentOffsets<P> {
magic: 0,
builtins: 0,
vm_store_context: 0,
flags: 0,
may_leave: 0,
task_may_block: 0,
trampoline_func_refs: 0,
intrinsic_func_refs: 0,
Expand Down Expand Up @@ -176,7 +172,7 @@ impl<P: PtrSize> VMComponentOffsets<P> {
size(builtins) = ret.ptr.size(),
size(vm_store_context) = ret.ptr.size(),
align(16),
size(flags) = cmul(ret.num_runtime_component_instances, ret.ptr.size_of_vmglobal_definition()),
size(may_leave) = cmul(ret.num_runtime_component_instances, ret.ptr.size_of_vmglobal_definition()),
size(task_may_block) = ret.ptr.size_of_vmglobal_definition(),
align(u32::from(ret.ptr.size())),
size(trampoline_func_refs) = cmul(ret.num_trampolines, ret.ptr.size_of_vm_func_ref()),
Expand Down Expand Up @@ -218,11 +214,11 @@ impl<P: PtrSize> VMComponentOffsets<P> {
self.builtins
}

/// The offset of the `flags` field.
/// The offset of the `may_leave` flag for the given component instance.
#[inline]
pub fn instance_flags(&self, index: RuntimeComponentInstanceIndex) -> u32 {
pub fn may_leave(&self, index: RuntimeComponentInstanceIndex) -> u32 {
assert!(index.as_u32() < self.num_runtime_component_instances);
self.flags + index.as_u32() * u32::from(self.ptr.size_of_vmglobal_definition())
self.may_leave + index.as_u32() * u32::from(self.ptr.size_of_vmglobal_definition())
}

/// The offset of the `task_may_block` field.
Expand Down
129 changes: 90 additions & 39 deletions crates/environ/src/fact/trampoline.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,13 +16,12 @@
//! can be somewhat arbitrary, an intentional decision.

use crate::component::{
CanonicalAbiInfo, ComponentTypesBuilder, FLAG_MAY_LEAVE, FixedEncoding as FE, FlatType,
InterfaceType, MAX_FLAT_ASYNC_PARAMS, MAX_FLAT_PARAMS, PREPARE_ASYNC_NO_RESULT,
PREPARE_ASYNC_WITH_RESULT, START_FLAG_ASYNC_CALLEE, StringEncoding, Transcode,
TypeComponentLocalErrorContextTableIndex, TypeEnumIndex, TypeFixedLengthListIndex,
TypeFlagsIndex, TypeFutureTableIndex, TypeListIndex, TypeMapIndex, TypeOptionIndex,
TypeRecordIndex, TypeResourceTableIndex, TypeResultIndex, TypeStreamTableIndex, TypeTupleIndex,
TypeVariantIndex, VariantInfo,
CanonicalAbiInfo, ComponentTypesBuilder, FixedEncoding as FE, FlatType, InterfaceType,
MAX_FLAT_ASYNC_PARAMS, MAX_FLAT_PARAMS, PREPARE_ASYNC_NO_RESULT, PREPARE_ASYNC_WITH_RESULT,
START_FLAG_ASYNC_CALLEE, StringEncoding, Transcode, TypeComponentLocalErrorContextTableIndex,
TypeEnumIndex, TypeFixedLengthListIndex, TypeFlagsIndex, TypeFutureTableIndex, TypeListIndex,
TypeMapIndex, TypeOptionIndex, TypeRecordIndex, TypeResourceTableIndex, TypeResultIndex,
TypeStreamTableIndex, TypeTupleIndex, TypeVariantIndex, VariantInfo,
};
use crate::fact::signature::Signature;
use crate::fact::transcode::Transcoder;
Expand Down Expand Up @@ -701,9 +700,9 @@ impl<'a, 'b> Compiler<'a, 'b> {
.map(|(i, ty)| (i as u32, *ty))
.collect::<Vec<_>>();

self.set_flag(adapter.lift.flags, FLAG_MAY_LEAVE, false);
let saved = self.clear_may_leave(adapter.lift.flags);
self.translate_params(adapter, &param_locals);
self.set_flag(adapter.lift.flags, FLAG_MAY_LEAVE, true);
self.restore_may_leave(adapter.lift.flags, saved);

self.finish();
}
Expand All @@ -727,7 +726,7 @@ impl<'a, 'b> Compiler<'a, 'b> {
.map(|(i, ty)| (i as u32, *ty))
.collect::<Vec<_>>();

self.set_flag(adapter.lower.flags, FLAG_MAY_LEAVE, false);
let saved = self.clear_may_leave(adapter.lower.flags);
// Note that we pass `param_locals` as _both_ the `param_locals` and
// `result_locals` parameters to `translate_results`. That's because
// the _parameters_ to `task.return` are actually the _results_ that the
Expand All @@ -739,7 +738,7 @@ impl<'a, 'b> Compiler<'a, 'b> {
// the import is lowered async, in which case `translate_results` will
// use that pointer to store the results.
self.translate_results(adapter, &param_locals, &param_locals);
self.set_flag(adapter.lower.flags, FLAG_MAY_LEAVE, true);
self.restore_may_leave(adapter.lower.flags, saved);

self.finish()
}
Expand All @@ -763,11 +762,12 @@ impl<'a, 'b> Compiler<'a, 'b> {
// This inserts the initial check required by `canon_lower` that the
// caller instance can be left and additionally checks the
// flags on the callee if necessary whether it can be entered.
self.trap_if_not_flag(
adapter.lower.flags,
FLAG_MAY_LEAVE,
Trap::CannotLeaveComponent,
);
//
// The loaded `may_leave` value is saved into `saved_lower_may_leave`
// so that it can be restored after results are translated below
// without reloading the global.
let saved_lower_may_leave =
self.trap_if_not_may_leave(adapter.lower.flags, Trap::CannotLeaveComponent);

let old_task_may_block = if self.module.tunables.concurrency_support {
// Save, clear, and later restore the `may_block` field.
Expand Down Expand Up @@ -823,27 +823,49 @@ impl<'a, 'b> Compiler<'a, 'b> {
None
};

// Perform the translation of arguments. Note that `FLAG_MAY_LEAVE` is
// cleared around this invocation for the callee as per the
// Perform the translation of arguments. Note that the `may_leave` flag
// is cleared around this invocation for the callee as per the
// `canon_lift` definition in the spec. Additionally note that the
// precise ordering of traps here is not required since internal state
// is not visible to either instance and a trap will "lock down" both
// instances to no longer be visible. This means that we're free to
// reorder lifts/lowers and flags and such as is necessary and
// convenient here.
//
// TODO: if translation doesn't actually call any functions in either
// instance then there's no need to set/clear the flag here and that can
// be optimized away.
self.set_flag(adapter.lift.flags, FLAG_MAY_LEAVE, false);
// The clear-and-restore is structured (a constant `0` store to clear,
// then a store of the saved original value to restore) so that if
// translation doesn't actually call any functions in either instance
// then a future dead-store elimination pass in Cranelift can remove all
// the flag juggling entirely (other than trapping when `!may_leave`):
//
// may_leave = load vmctx+MAY_LEAVE_OFFSET ;; (0)
// trapz may_leave
//
// ...
//
// zero = iconst 0
// store zero, vmctx+MAY_LEAVE_OFFSET ;; (1)
//
// ...
//
// store may_leave, vmctx+MAY_LEAVE_OFFSET ;; (2)
//
// First, the dead-store elimination pass will see that the the store at
// (1) is dead and remove it. Then, the idempotent-store eliminator will
// recognize that the store at (2) is storing the same value that the
// memory location already contains and it will also be removed. The
// more we can reuse locals to make this idempotency obvious, rather
// than force Cranelift's optimizer to rediscover this information, the
// better.
let saved_lift_may_leave = self.clear_may_leave(adapter.lift.flags);
let param_locals = lower_sig
.params
.iter()
.enumerate()
.map(|(i, ty)| (i as u32, *ty))
.collect::<Vec<_>>();
self.translate_params(adapter, &param_locals);
self.set_flag(adapter.lift.flags, FLAG_MAY_LEAVE, true);
self.restore_may_leave(adapter.lift.flags, saved_lift_may_leave);

// With all the arguments on the stack the actual target function is
// now invoked. The core wasm results of the function are then placed
Expand Down Expand Up @@ -882,12 +904,9 @@ impl<'a, 'b> Compiler<'a, 'b> {
// order of everything doesn't matter since intermediate states cannot
// be witnessed, hence the setting of flags here to encapsulate both
// liftings and lowerings.
//
// TODO: like above the management of the `MAY_LEAVE` flag can probably
// be elided here for "simple" results.
self.set_flag(adapter.lower.flags, FLAG_MAY_LEAVE, false);
self.set_may_leave_false(adapter.lower.flags);
self.translate_results(adapter, &param_locals, &result_locals);
self.set_flag(adapter.lower.flags, FLAG_MAY_LEAVE, true);
self.restore_may_leave(adapter.lower.flags, saved_lower_may_leave);

// And finally post-return state is handled here once all results/etc
// are all translated.
Expand Down Expand Up @@ -3575,26 +3594,58 @@ impl<'a, 'b> Compiler<'a, 'b> {
}
}

fn trap_if_not_flag(&mut self, flags_global: GlobalIndex, flag_to_test: i32, trap: Trap) {
/// Loads the `may_leave` flag for the given instance, traps with `trap` if
/// it is not set, and returns a temporary local holding the loaded value
/// so that it can later be restored with `restore_may_leave` without
/// reloading the global.
///
/// The `may_leave` flag is a boolean (0 or 1) so no masking is required.
fn trap_if_not_may_leave(&mut self, flags_global: GlobalIndex, trap: Trap) -> TempLocal {
self.instruction(GlobalGet(flags_global.as_u32()));
self.instruction(I32Const(flag_to_test));
self.instruction(I32And);
// Save the flag's value (known to be `true` whenever the trap below is
// not taken) into a temporary for later restoration.
let saved = self.local_tee_new_tmp(ValType::I32);
self.instruction(I32Eqz);
self.instruction(If(BlockType::Empty));
self.trap(trap);
self.instruction(End);
saved
}

fn set_flag(&mut self, flags_global: GlobalIndex, flag_to_set: i32, value: bool) {
/// Saves the current value of the `may_leave` flag into a fresh temporary
/// local (returned) and then clears the flag to `false`.
fn clear_may_leave(&mut self, flags_global: GlobalIndex) -> TempLocal {
self.instruction(GlobalGet(flags_global.as_u32()));
if value {
self.instruction(I32Const(flag_to_set));
self.instruction(I32Or);
} else {
self.instruction(I32Const(!flag_to_set));
self.instruction(I32And);
}
let saved = self.local_set_new_tmp(ValType::I32);
self.set_may_leave_false(flags_global);
saved
}

/// Sets the `may_leave` flag to `false` by storing a constant `0`.
///
/// Since there is only a single flag there is no need to reload the global
/// and mask: storing `0` is sufficient.
fn set_may_leave_false(&mut self, flags_global: GlobalIndex) {
self.instruction(I32Const(0));
self.instruction(GlobalSet(flags_global.as_u32()));
}

/// Restores the `may_leave` flag to the value previously saved in `saved`
/// (via `clear_may_leave` or `trap_if_not_may_leave`) and frees the
/// temporary local.
///
/// Storing the previously-loaded value (rather than reloading the global
/// and or-ing in the flag bit) makes it clear in the generated CLIF that
/// the same value that was there before is being written back. Combined
/// with the constant `0` store in `set_may_leave_false`, this lets a future
/// dead-store-elimination pass remove the clear-to-`false` store which will
/// then allow our idempotent-store elimination to remove this restore for
/// adapters whose body never touches the flag (e.g. simple inlined
/// callees).
fn restore_may_leave(&mut self, flags_global: GlobalIndex, saved: TempLocal) {
self.instruction(LocalGet(saved.idx));
self.instruction(GlobalSet(flags_global.as_u32()));
self.free_temp_local(saved);
}

fn assert_aligned(&mut self, ty: &InterfaceType, mem: &Memory) {
Expand Down
16 changes: 6 additions & 10 deletions crates/wasmtime/src/runtime/vm/component.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -365,8 +365,8 @@ impl ComponentInstance {
#[inline]
pub fn instance_flags(&self, instance: RuntimeComponentInstanceIndex) -> InstanceFlags {
unsafe {
let ptr = self
.vmctx_plus_offset_raw::<VMGlobalDefinition>(self.offsets.instance_flags(instance));
let ptr =
self.vmctx_plus_offset_raw::<VMGlobalDefinition>(self.offsets.may_leave(instance));
InstanceFlags(SendSyncPtr::new(ptr))
}
}
Expand Down Expand Up @@ -704,9 +704,9 @@ impl ComponentInstance {
let i = RuntimeComponentInstanceIndex::from_u32(i);
let mut def = VMGlobalDefinition::new();
// SAFETY: this is a valid initialization of all globals which are
// 32-bit values.
// 32-bit values. The `may_leave` flag starts out `true`.
unsafe {
*def.as_i32_mut() = FLAG_MAY_LEAVE;
*def.as_i32_mut() = 1;
self.instance_flags(i).as_raw().write(def);
}
}
Expand Down Expand Up @@ -1062,17 +1062,13 @@ impl InstanceFlags {

#[inline]
pub unsafe fn may_leave(&self) -> bool {
unsafe { *self.as_raw().as_ref().as_i32() & FLAG_MAY_LEAVE != 0 }
unsafe { *self.as_raw().as_ref().as_i32() != 0 }
}

#[inline]
pub unsafe fn set_may_leave(&mut self, val: bool) {
unsafe {
if val {
*self.as_raw().as_mut().as_i32_mut() |= FLAG_MAY_LEAVE;
} else {
*self.as_raw().as_mut().as_i32_mut() &= !FLAG_MAY_LEAVE;
}
*self.as_raw().as_mut().as_i32_mut() = val as i32;
}
}

Expand Down
Loading
Loading