Inline context.{get,set} in components

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,70 @@ 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) {
+        let ty = match intrinsic {
+            UnsafeIntrinsic::ContextGetI32_0
+            | UnsafeIntrinsic::ContextSetI32_0
+            | UnsafeIntrinsic::ContextGetI32_1
+            | UnsafeIntrinsic::ContextSetI32_1 => ir::types::I32,
+            _ => unreachable!(),
+        };
+        let context_slot_size = 4;
+        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 * context_slot_size;
+        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 +1673,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 +1751,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 +1776,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

@@ -250,6 +250,16 @@ 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` field of `VMStoreContext`.
+    fn vmstore_context_component_context(&self) -> u8 {
+        self.vmstore_context_async_guard_range() + 2 * self.size()
+    }
+
     // Offsets within `VMMemoryDefinition`
 
     /// The offset of the `base` field.