Inline context.{get,set} in components (#13194)

* Inline `context.{get,set}` in components

This commit reimplements the `context.{get,set}` intrinsics in the
component model, introduced in the component-model-async and
component-model-threading proposals. The intent of these intrinsics in
WASIp3, for example, are intended to replace the `global`s used for the
stack pointer and TLS base in previous modules, for example. The
implementation of loading from a `global` is a single load instruction,
whereas the previous implementation of `context.get` was a full libcall,
which is significantly more expensive. The goal of this PR is to ensure
that the transition to using `context.get` and `context.set` for
high-performance uses retains the same performance as the WASIp2
constructs.

Specifically the storage for `context.{get,set}` slots have been moved
into the `VMStoreContext` structure which has a known layout to compiled
code. There still remains storage within each `GuestThread` because
there's only one store, and the idea is that whenever threads are
switched between the switch operation is slightly more expensive now
where it has to update and maintain the state in the store. The
rationale for this is that it'll be far more often that these values are
accessed rather than threads being swapped between.

The implementation chosen in this commit is to model the
`context.{get,set}` intrinsics as `UnsafeIntrinsic`s. This is a bit of a
shoehorn where they're not actually unsafe, but all of the plumbing and
support for `UnsafeIntrinsic` is effectively exactly what these want. To
avoid duplicating lots of infrastructure that's where these now reside.

The `concurrent.rs` implementation has been updated to save/restore
context from the store, and this additionally updates a few other switch
points to ensure that the store never switches away or to a deleted
thread. This niche situation happened in a few scenarios with no impact
from before, but with the switching implementation having to access
threads it became load-bearing that these must be valid.

The end result is that with `-Cinlining` the `context.{get,set}`
instructions are two instructions instead of a full libcall. One
instruction is loading `VMStoreContext`, which is GVN-able and
hoist-able, while the other is the actual load/store. This is the same
as the performance of the stack pointer being in an imported global,
for example.

* Fix build issues

* Review comments

Author: alexcrichton

crates/cranelift/src/compiler/component.rs

@@ -724,28 +724,6 @@ impl<'a> TrampolineCompiler<'a> {
                     |_, _| {},
                 );
             }
-            Trampoline::ContextGet { instance, slot } => {
-                self.translate_libcall(
-                    host::context_get,
-                    TrapSentinel::NegativeOne,
-                    WasmArgs::InRegisters,
-                    |me, params| {
-                        params.push(me.index_value(*instance));
-                        params.push(me.builder.ins().iconst(ir::types::I32, i64::from(*slot)));
-                    },
-                );
-            }
-            Trampoline::ContextSet { instance, slot } => {
-                self.translate_libcall(
-                    host::context_set,
-                    TrapSentinel::Falsy,
-                    WasmArgs::InRegisters,
-                    |me, params| {
-                        params.push(me.index_value(*instance));
-                        params.push(me.builder.ins().iconst(ir::types::I32, i64::from(*slot)));
-                    },
-                );
-            }
             Trampoline::ThreadIndex => {
                 self.translate_libcall(
                     host::thread_index,
@@ -1464,9 +1442,7 @@ impl<'a> TrampolineCompiler<'a> {
             Trampoline::ResourceRep { .. }
             | Trampoline::ThreadIndex
             | Trampoline::BackpressureInc { .. }
-            | Trampoline::BackpressureDec { .. }
-            | Trampoline::ContextGet { .. }
-            | Trampoline::ContextSet { .. } => return,
+            | Trampoline::BackpressureDec { .. } => return,
 
             // Intrinsics used in adapters generated by FACT that aren't called
             // directly from guest wasm, so no check is needed.
@@ -1552,6 +1528,74 @@ impl<'a> TrampolineCompiler<'a> {
             &mut self.builder,
         )
     }
+
+    /// Loads `*mut VMStoreContext` and returns it.
+    ///
+    /// Note that the `*mut VMStoreContext` value is the same for all
+    /// `VMContext`-like structures in a store. In this case it's loaded from
+    /// the *caller* vmctx rather than the *callee* vmctx. The caller is using a
+    /// `VMContext` for core wasm which is passed in a register, where the
+    /// callee is a `VMComponentContext` loaded from the `VMContext`. By using
+    /// the caller vmctx we're able to possibly eliminate the dead load of the
+    /// `VMComponentContext` if it's otherwise unused.
+    fn load_vm_store_context(&mut self) -> ir::Value {
+        let caller_vmctx = self.abi_load_params()[1];
+        self.builder.ins().load(
+            self.isa.pointer_type(),
+            ir::MemFlags::trusted()
+                .with_readonly()
+                .with_alias_region(Some(ir::AliasRegion::Vmctx))
+                .with_can_move(),
+            caller_vmctx,
+            i32::from(self.offsets.ptr.vmctx_store_context()),
+        )
+    }
+
+    fn translate_context_intrinsic(&mut self, intrinsic: UnsafeIntrinsic) {
+        // This is the width of the type being loaded from Wasmtime's
+        // `VMStoreContext` slot and it depends on the intrinsic.
+        let ty = match intrinsic {
+            UnsafeIntrinsic::ContextGetI32_0
+            | UnsafeIntrinsic::ContextSetI32_0
+            | UnsafeIntrinsic::ContextGetI32_1
+            | UnsafeIntrinsic::ContextSetI32_1 => ir::types::I32,
+            _ => unreachable!(),
+        };
+
+        let slot = match intrinsic {
+            UnsafeIntrinsic::ContextGetI32_0 | UnsafeIntrinsic::ContextSetI32_0 => 0,
+            UnsafeIntrinsic::ContextGetI32_1 | UnsafeIntrinsic::ContextSetI32_1 => 1,
+            _ => unreachable!(),
+        };
+        let offset = self
+            .offsets
+            .ptr
+            .vmstore_context_component_context_slot(slot);
+        let params = self.abi_load_params();
+        let vmstore_context = self.load_vm_store_context();
+        match intrinsic {
+            UnsafeIntrinsic::ContextGetI32_0 | UnsafeIntrinsic::ContextGetI32_1 => {
+                let context = self.builder.ins().load(
+                    ty,
+                    MemFlags::trusted(),
+                    vmstore_context,
+                    i32::from(offset),
+                );
+                self.abi_store_results(&[context]);
+            }
+            UnsafeIntrinsic::ContextSetI32_0 | UnsafeIntrinsic::ContextSetI32_1 => {
+                let new_context = params[2];
+                self.builder.ins().store(
+                    MemFlags::trusted(),
+                    new_context,
+                    vmstore_context,
+                    i32::from(offset),
+                );
+                self.abi_store_results(&[]);
+            }
+            _ => unreachable!(),
+        }
+    }
 }
 
 // Helper structure to implement `TranslateTrap`. This isn't possible to do
@@ -1633,12 +1677,7 @@ impl ComponentCompiler for Compiler {
             vmctx,
             wasmtime_environ::component::VMCOMPONENT_MAGIC,
         );
-        let vm_store_context = c.builder.ins().load(
-            pointer_type,
-            MemFlags::trusted(),
-            vmctx,
-            i32::try_from(c.offsets.vm_store_context()).unwrap(),
-        );
+        let vm_store_context = c.load_vm_store_context();
         super::save_last_wasm_exit_fp_and_pc(
             &mut c.builder,
             pointer_type,
@@ -1716,21 +1755,10 @@ impl ComponentCompiler for Compiler {
             | UnsafeIntrinsic::U32NativeStore
             | UnsafeIntrinsic::U64NativeStore => c.translate_store_intrinsic(intrinsic)?,
             UnsafeIntrinsic::StoreDataAddress => {
-                let [callee_vmctx, _caller_vmctx] = *c.abi_load_params() else {
-                    unreachable!()
-                };
                 let pointer_type = self.isa.pointer_type();
 
                 // Load the `*mut VMStoreContext` out of our vmctx.
-                let store_ctx = c.builder.ins().load(
-                    pointer_type,
-                    ir::MemFlags::trusted()
-                        .with_readonly()
-                        .with_alias_region(Some(ir::AliasRegion::Vmctx))
-                        .with_can_move(),
-                    callee_vmctx,
-                    i32::try_from(c.offsets.vm_store_context()).unwrap(),
-                );
+                let store_ctx = c.load_vm_store_context();
 
                 // Load the `*mut T` out of the `VMStoreContext`.
                 let data_address = c.builder.ins().load(
@@ -1752,6 +1780,13 @@ impl ComponentCompiler for Compiler {
 
                 c.abi_store_results(&[data_address]);
             }
+
+            UnsafeIntrinsic::ContextGetI32_0
+            | UnsafeIntrinsic::ContextGetI32_1
+            | UnsafeIntrinsic::ContextSetI32_0
+            | UnsafeIntrinsic::ContextSetI32_1 => {
+                c.translate_context_intrinsic(intrinsic);
+            }
         }
 
         c.builder.finalize();

crates/environ/src/component.rs

@@ -186,10 +186,6 @@ macro_rules! foreach_builtin_component_function {
             #[cfg(feature = "component-model-async")]
             error_context_transfer(vmctx: vmctx, src_idx: u32, src_table: u32, dst_table: u32) -> u64;
             #[cfg(feature = "component-model-async")]
-            context_get(vmctx: vmctx, caller_instance: u32, slot: u32) -> u64;
-            #[cfg(feature = "component-model-async")]
-            context_set(vmctx: vmctx, caller_instance: u32, slot: u32, val: u32) -> bool;
-            #[cfg(feature = "component-model-async")]
             thread_index(vmctx: vmctx) -> u64;
             #[cfg(feature = "component-model-async")]
             thread_new_indirect(vmctx: vmctx, caller_instance: u32, func_ty_id: u32, func_table_idx: u32, func_idx: u32, context: u32) -> u64;

crates/environ/src/component/dfg.rs

@@ -476,14 +476,6 @@ pub enum Trampoline {
     Trap,
     EnterSyncCall,
     ExitSyncCall,
-    ContextGet {
-        instance: RuntimeComponentInstanceIndex,
-        slot: u32,
-    },
-    ContextSet {
-        instance: RuntimeComponentInstanceIndex,
-        slot: u32,
-    },
     ThreadIndex,
     ThreadNewIndirect {
         instance: RuntimeComponentInstanceIndex,
@@ -1167,14 +1159,6 @@ impl LinearizeDfg<'_> {
             Trampoline::Trap => info::Trampoline::Trap,
             Trampoline::EnterSyncCall => info::Trampoline::EnterSyncCall,
             Trampoline::ExitSyncCall => info::Trampoline::ExitSyncCall,
-            Trampoline::ContextGet { instance, slot } => info::Trampoline::ContextGet {
-                instance: *instance,
-                slot: *slot,
-            },
-            Trampoline::ContextSet { instance, slot } => info::Trampoline::ContextSet {
-                instance: *instance,
-                slot: *slot,
-            },
             Trampoline::ThreadIndex => info::Trampoline::ThreadIndex,
             Trampoline::ThreadNewIndirect {
                 instance,

crates/environ/src/component/info.rs

@@ -1118,28 +1118,6 @@ pub enum Trampoline {
     /// pushed by `EnterSyncCall`.
     ExitSyncCall,
 
-    /// Intrinsic used to implement the `context.get` component model builtin.
-    ///
-    /// The payload here represents that this is accessing the Nth slot of local
-    /// storage.
-    ContextGet {
-        /// The specific component instance which is calling the intrinsic.
-        instance: RuntimeComponentInstanceIndex,
-        /// Which slot to access.
-        slot: u32,
-    },
-
-    /// Intrinsic used to implement the `context.set` component model builtin.
-    ///
-    /// The payload here represents that this is accessing the Nth slot of local
-    /// storage.
-    ContextSet {
-        /// The specific component instance which is calling the intrinsic.
-        instance: RuntimeComponentInstanceIndex,
-        /// Which slot to update.
-        slot: u32,
-    },
-
     /// Intrinsic used to implement the `thread.index` component model builtin.
     ThreadIndex,
 
@@ -1256,8 +1234,6 @@ impl Trampoline {
             Trap => format!("trap"),
             EnterSyncCall => format!("enter-sync-call"),
             ExitSyncCall => format!("exit-sync-call"),
-            ContextGet { .. } => format!("context-get"),
-            ContextSet { .. } => format!("context-set"),
             ThreadIndex => format!("thread-index"),
             ThreadNewIndirect { .. } => format!("thread-new-indirect"),
             ThreadSuspendToSuspended { .. } => format!("thread-suspend-to-suspended"),

crates/environ/src/component/intrinsic.rs

@@ -22,6 +22,11 @@ macro_rules! for_each_unsafe_intrinsic {
 
             "u64-native-load" => U64NativeLoad : u64_native_load(address: u64) -> u64;
             "u64-native-store" => U64NativeStore : u64_native_store(address: u64, value: u64);
+
+            "context-get-i32-0" => ContextGetI32_0 : context_get_i32_0() -> u32;
+            "context-set-i32-0" => ContextSetI32_0 : context_set_i32_0(val: u32);
+            "context-get-i32-1" => ContextGetI32_1 : context_get_i32_1() -> u32;
+            "context-set-i32-1" => ContextSetI32_1 : context_set_i32_1(val: u32);
         }
     };
 }

crates/environ/src/component/translate/inline.rs

@@ -1074,24 +1074,24 @@ impl<'a> Inliner<'a> {
                 frame.funcs.push((*func, dfg::CoreDef::Trampoline(index)));
             }
             ContextGet { func, i } => {
-                let index = self.result.trampolines.push((
-                    *func,
-                    dfg::Trampoline::ContextGet {
-                        instance: frame.instance,
-                        slot: *i,
-                    },
-                ));
-                frame.funcs.push((*func, dfg::CoreDef::Trampoline(index)));
+                let intrinsic = match i {
+                    0 => UnsafeIntrinsic::ContextGetI32_0,
+                    1 => UnsafeIntrinsic::ContextGetI32_1,
+                    _ => unreachable!(),
+                };
+                frame
+                    .funcs
+                    .push((*func, dfg::CoreDef::UnsafeIntrinsic(*func, intrinsic)));
             }
             ContextSet { func, i } => {
-                let index = self.result.trampolines.push((
-                    *func,
-                    dfg::Trampoline::ContextSet {
-                        instance: frame.instance,
-                        slot: *i,
-                    },
-                ));
-                frame.funcs.push((*func, dfg::CoreDef::Trampoline(index)));
+                let intrinsic = match i {
+                    0 => UnsafeIntrinsic::ContextSetI32_0,
+                    1 => UnsafeIntrinsic::ContextSetI32_1,
+                    _ => unreachable!(),
+                };
+                frame
+                    .funcs
+                    .push((*func, dfg::CoreDef::UnsafeIntrinsic(*func, intrinsic)));
             }
             ThreadIndex { func } => {
                 let index = self

crates/environ/src/vmoffsets.rs

@@ -39,6 +39,10 @@ use crate::{
 };
 use cranelift_entity::packed_option::ReservedValue;
 
+/// Number of slots in for `component_context` in the `VMStoreContext`. This is
+/// defined by the component model's `context.{get,set}` intrinsics.
+pub const NUM_COMPONENT_CONTEXT_SLOTS: usize = 2;
+
 #[cfg(target_pointer_width = "32")]
 fn cast_to_u32(sz: usize) -> u32 {
     u32::try_from(sz).unwrap()
@@ -250,6 +254,20 @@ pub trait PtrSize {
         self.vmstore_context_stack_chain() + self.size_of_vmstack_chain()
     }
 
+    /// Return the offset of the `async_guard_range` field of `VMStoreContext`.
+    fn vmstore_context_async_guard_range(&self) -> u8 {
+        self.vmstore_context_store_data() + self.size()
+    }
+
+    /// Return the offset of the `component_context[i]` field of
+    /// `VMStoreContext`.
+    fn vmstore_context_component_context_slot(&self, i: u8) -> u8 {
+        assert!(usize::from(i) < NUM_COMPONENT_CONTEXT_SLOTS);
+        let base = self.vmstore_context_async_guard_range() + 2 * self.size();
+        let slot_size = 4;
+        base + i * slot_size
+    }
+
     // Offsets within `VMMemoryDefinition`
 
     /// The offset of the `base` field.