Repository: bytecodealliance/wasmtime
Pull request #8190 (2024-03-20 15:02:46)

Implement opt-in for enabling WASI to block the current thread
Currently all of Wasmtime's implementation of WASI is built on Tokio, but some operations are currently not asynchronous such as opening a file or reading a directory. Internally these use `spawn_blocking` to satisfy the requirements of async users of WASI to avoid blocking the current thread. This use of `spawn_blocking`, however, ends up causing mostly just performance overhead in the use case of the CLI, for example, where async wasm is not used. That then leads to this commit, implementing an opt-in mechanism to be able to block the current thread.

A `WasiCtx` now has a flag indicating whether it's ok to block the current thread and that's carried to various filesystem operations that use `spawn_blocking`. The call to `spawn_blocking` is now conditional and skipped if this flag is set.

Additionally the invocation point in the CLI for wasm modules is wrapped in a Tokio runtime to avoid entering/exiting Tokio in the "leaves" when wasm calls the host, as happens today. This hits a better fast path in Tokio that appears to be more efficient.

Semantically this should not result in any change for CLI programs except in one case: file writes. By default writes on `output-stream` in WASI are asynchronous meaning that only one write can be in flight at a time. That being said all current users are immediately blocking waiting for this write to finish anyway, so this commit won't end up changing much. It's already the case that file reads are always blocking, for example. If necessary in the future though this can be further special-cased at the preview1 layer.

<!--
Please make sure you include the following information:

- If this work has been discussed elsewhere, please include a link to that
  conversation. If it was discussed in an issue, just mention "issue #...".

- Explain why this change is needed. If the details are in an issue already,
  this can be brief.

Our development process is documented in the Wasmtime book:
https://docs.wasmtime.dev/contributing-development-process.html

Please ensure all communication follows the code of conduct:
https://github.com/bytecodealliance/wasmtime/blob/main/CODE_OF_CONDUCT.md
-->



Patch:
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index e606cf6bb..f671b6176 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -340,7 +340,7 @@ jobs:
     # Check whether `wasmtime` cross-compiles to aarch64-linux-android
     - run: rustup target add aarch64-linux-android
     - name: Set up JDK 17
-      uses: actions/setup-java@v3
+      uses: actions/setup-java@v4
       with:
         java-version: '17'
         distribution: 'temurin'
diff --git a/build.rs b/build.rs
index 262b40838..d0b42afe2 100644
--- a/build.rs
+++ b/build.rs
@@ -241,6 +241,10 @@ fn ignore(testsuite: &str, testname: &str, strategy: &str) -> bool {
             return denylist || ref_types || simd;
         }
 
+        if testsuite == "memory64" {
+            return testname.starts_with("simd") || testname.starts_with("threads");
+        }
+
         if testsuite != "winch" {
             return true;
         }
diff --git a/cranelift/codegen/src/isa/aarch64/abi.rs b/cranelift/codegen/src/isa/aarch64/abi.rs
index 440840c4d..31a5c25c3 100644
--- a/cranelift/codegen/src/isa/aarch64/abi.rs
+++ b/cranelift/codegen/src/isa/aarch64/abi.rs
@@ -99,17 +99,14 @@ impl ABIMachineSpec for AArch64MachineDeps {
         16
     }
 
-    fn compute_arg_locs<'a, I>(
+    fn compute_arg_locs(
         call_conv: isa::CallConv,
         _flags: &settings::Flags,
-        params: I,
+        params: &[ir::AbiParam],
         args_or_rets: ArgsOrRets,
         add_ret_area_ptr: bool,
-        mut args: ArgsAccumulator<'_>,
-    ) -> CodegenResult<(u32, Option<usize>)>
-    where
-        I: IntoIterator<Item = &'a ir::AbiParam>,
-    {
+        mut args: ArgsAccumulator,
+    ) -> CodegenResult<(u32, Option<usize>)> {
         if matches!(call_conv, isa::CallConv::Tail | isa::CallConv::Winch) {
             return compute_arg_locs_tail(params, add_ret_area_ptr, args);
         }
@@ -1291,14 +1288,11 @@ impl AArch64CallSite {
     }
 }
 
-fn compute_arg_locs_tail<'a, I>(
-    params: I,
+fn compute_arg_locs_tail(
+    params: &[ir::AbiParam],
     add_ret_area_ptr: bool,
-    mut args: ArgsAccumulator<'_>,
-) -> CodegenResult<(u32, Option<usize>)>
-where
-    I: IntoIterator<Item = &'a ir::AbiParam>,
-{
+    mut args: ArgsAccumulator,
+) -> CodegenResult<(u32, Option<usize>)> {
     let mut xregs = ALL_CLOBBERS
         .into_iter()
         .filter(|r| r.class() == RegClass::Int)
diff --git a/cranelift/codegen/src/isa/riscv64/abi.rs b/cranelift/codegen/src/isa/riscv64/abi.rs
index 0f0d32e8f..c119fe312 100644
--- a/cranelift/codegen/src/isa/riscv64/abi.rs
+++ b/cranelift/codegen/src/isa/riscv64/abi.rs
@@ -87,17 +87,14 @@ impl ABIMachineSpec for Riscv64MachineDeps {
         16
     }
 
-    fn compute_arg_locs<'a, I>(
+    fn compute_arg_locs(
         call_conv: isa::CallConv,
         _flags: &settings::Flags,
-        params: I,
+        params: &[ir::AbiParam],
         args_or_rets: ArgsOrRets,
         add_ret_area_ptr: bool,
-        mut args: ArgsAccumulator<'_>,
-    ) -> CodegenResult<(u32, Option<usize>)>
-    where
-        I: IntoIterator<Item = &'a ir::AbiParam>,
-    {
+        mut args: ArgsAccumulator,
+    ) -> CodegenResult<(u32, Option<usize>)> {
         assert_ne!(
             call_conv,
             isa::CallConv::Winch,
diff --git a/cranelift/codegen/src/isa/s390x/abi.rs b/cranelift/codegen/src/isa/s390x/abi.rs
index af34bb5a4..7c48578a8 100644
--- a/cranelift/codegen/src/isa/s390x/abi.rs
+++ b/cranelift/codegen/src/isa/s390x/abi.rs
@@ -220,17 +220,14 @@ impl ABIMachineSpec for S390xMachineDeps {
         8
     }
 
-    fn compute_arg_locs<'a, I>(
+    fn compute_arg_locs(
         call_conv: isa::CallConv,
         _flags: &settings::Flags,
-        params: I,
+        params: &[ir::AbiParam],
         args_or_rets: ArgsOrRets,
         add_ret_area_ptr: bool,
-        mut args: ArgsAccumulator<'_>,
-    ) -> CodegenResult<(u32, Option<usize>)>
-    where
-        I: IntoIterator<Item = &'a ir::AbiParam>,
-    {
+        mut args: ArgsAccumulator,
+    ) -> CodegenResult<(u32, Option<usize>)> {
         assert_ne!(
             call_conv,
             isa::CallConv::Tail,
diff --git a/cranelift/codegen/src/isa/x64/abi.rs b/cranelift/codegen/src/isa/x64/abi.rs
index eeda5b6cd..9b8eb0764 100644
--- a/cranelift/codegen/src/isa/x64/abi.rs
+++ b/cranelift/codegen/src/isa/x64/abi.rs
@@ -101,17 +101,14 @@ impl ABIMachineSpec for X64ABIMachineSpec {
         16
     }
 
-    fn compute_arg_locs<'a, I>(
+    fn compute_arg_locs(
         call_conv: isa::CallConv,
         flags: &settings::Flags,
-        params: I,
+        params: &[ir::AbiParam],
         args_or_rets: ArgsOrRets,
         add_ret_area_ptr: bool,
-        mut args: ArgsAccumulator<'_>,
-    ) -> CodegenResult<(u32, Option<usize>)>
-    where
-        I: IntoIterator<Item = &'a ir::AbiParam>,
-    {
+        mut args: ArgsAccumulator,
+    ) -> CodegenResult<(u32, Option<usize>)> {
         let is_fastcall = call_conv.extends_windows_fastcall();
 
         let mut next_gpr = 0;
diff --git a/cranelift/codegen/src/machinst/abi.rs b/cranelift/codegen/src/machinst/abi.rs
index 8e16fdae8..4af7bc343 100644
--- a/cranelift/codegen/src/machinst/abi.rs
+++ b/cranelift/codegen/src/machinst/abi.rs
@@ -408,16 +408,14 @@ pub trait ABIMachineSpec {
     /// Returns the stack-space used (rounded up to as alignment requires), and
     /// if `add_ret_area_ptr` was passed, the index of the extra synthetic arg
     /// that was added.
-    fn compute_arg_locs<'a, I>(
+    fn compute_arg_locs(
         call_conv: isa::CallConv,
         flags: &settings::Flags,
-        params: I,
+        params: &[ir::AbiParam],
         args_or_rets: ArgsOrRets,
         add_ret_area_ptr: bool,
-        args: ArgsAccumulator<'_>,
-    ) -> CodegenResult<(u32, Option<usize>)>
-    where
-        I: IntoIterator<Item = &'a ir::AbiParam>;
+        args: ArgsAccumulator,
+    ) -> CodegenResult<(u32, Option<usize>)>;
 
     /// Returns the offset from FP to the argument area, i.e., jumping over the saved FP, return
     /// address, and maybe other standard elements depending on ABI (e.g. Wasm TLS reg).
@@ -827,8 +825,13 @@ impl SigSet {
         sig: &ir::Signature,
         flags: &settings::Flags,
     ) -> CodegenResult<SigData> {
-        let sret = missing_struct_return(sig);
-        let returns = sret.as_ref().into_iter().chain(&sig.returns);
+        use std::borrow::Cow;
+
+        let returns = if let Some(sret) = missing_struct_return(sig) {
+            Cow::from_iter(std::iter::once(&sret).chain(&sig.returns).copied())
+        } else {
+            Cow::from(sig.returns.as_slice())
+        };
 
         // Compute args and retvals from signature. Handle retvals first,
         // because we may need to add a return-area arg to the args.
@@ -839,7 +842,7 @@ impl SigSet {
         let (sized_stack_ret_space, _) = M::compute_arg_locs(
             sig.call_conv,
             flags,
-            returns,
+            &returns,
             ArgsOrRets::Rets,
             /* extra ret-area ptr = */ false,
             ArgsAccumulator::new(&mut self.abi_args),
diff --git a/crates/cranelift/src/func_environ.rs b/crates/cranelift/src/func_environ.rs
index 6104a9411..cf2b96f15 100644
--- a/crates/cranelift/src/func_environ.rs
+++ b/crates/cranelift/src/func_environ.rs
@@ -1037,6 +1037,87 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
             .get(&func_index)
             .map(|s| *s)
     }
+
+    /// Proof-carrying code: create a memtype describing an empty
+    /// runtime struct (to be updated later).
+    fn create_empty_struct_memtype(&self, func: &mut ir::Function) -> ir::MemoryType {
+        func.create_memory_type(ir::MemoryTypeData::Struct {
+            size: 0,
+            fields: vec![],
+        })
+    }
+
+    /// Proof-carrying code: add a new field to a memtype used to
+    /// describe a runtime struct. A memory region of type `memtype`
+    /// will have a pointer at `offset` pointing to another memory
+    /// region of type `pointee`. `readonly` indicates whether the
+    /// PCC-checked code is expected to update this field or not.
+    fn add_field_to_memtype(
+        &self,
+        func: &mut ir::Function,
+        memtype: ir::MemoryType,
+        offset: u32,
+        pointee: ir::MemoryType,
+        readonly: bool,
+    ) {
+        let ptr_size = self.pointer_type().bytes();
+        match &mut func.memory_types[memtype] {
+            ir::MemoryTypeData::Struct { size, fields } => {
+                *size = std::cmp::max(*size, offset.checked_add(ptr_size).unwrap().into());
+                fields.push(ir::MemoryTypeField {
+                    ty: self.pointer_type(),
+                    offset: offset.into(),
+                    readonly,
+                    fact: Some(ir::Fact::Mem {
+                        ty: pointee,
+                        min_offset: 0,
+                        max_offset: 0,
+                        nullable: false,
+                    }),
+                });
+
+                // Sort fields by offset -- we need to do this now
+                // because we may create an arbitrary number of
+                // memtypes for imported memories and we don't
+                // otherwise track them.
+                fields.sort_by_key(|f| f.offset);
+            }
+            _ => panic!("Cannot add field to non-struct memtype"),
+        }
+    }
+
+    /// Add one level of indirection to a pointer-and-memtype pair:
+    /// generate a load in the code at the specified offset, and if
+    /// memtypes are in use, add a field to the original struct and
+    /// generate a new memtype for the pointee.
+    fn load_pointer_with_memtypes(
+        &self,
+        func: &mut ir::Function,
+        value: ir::GlobalValue,
+        offset: u32,
+        readonly: bool,
+        memtype: Option<ir::MemoryType>,
+    ) -> (ir::GlobalValue, Option<ir::MemoryType>) {
+        let pointee = func.create_global_value(ir::GlobalValueData::Load {
+            base: value,
+            offset: Offset32::new(i32::try_from(offset).unwrap()),
+            global_type: self.pointer_type(),
+            flags: MemFlags::trusted().with_readonly(),
+        });
+
+        let mt = memtype.map(|mt| {
+            let pointee_mt = self.create_empty_struct_memtype(func);
+            self.add_field_to_memtype(func, mt, offset, pointee_mt, readonly);
+            func.global_value_facts[pointee] = Some(Fact::Mem {
+                ty: pointee_mt,
+                min_offset: 0,
+                max_offset: 0,
+                nullable: false,
+            });
+            pointee_mt
+        });
+        (pointee, mt)
+    }
 }
 
 struct Call<'a, 'func, 'module_env> {
@@ -2037,16 +2118,17 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
                     // VMMemoryDefinition` to it and dereference that when
                     // atomically growing it.
                     let from_offset = self.offsets.vmctx_vmmemory_pointer(def_index);
-                    let memory = func.create_global_value(ir::GlobalValueData::Load {
-                        base: vmctx,
-                        offset: Offset32::new(i32::try_from(from_offset).unwrap()),
-                        global_type: pointer_type,
-                        flags: MemFlags::trusted().with_readonly(),
-                    });
+                    let (memory, def_mt) = self.load_pointer_with_memtypes(
+                        func,
+                        vmctx,
+                        from_offset,
+                        true,
+                        self.pcc_vmctx_memtype,
+                    );
                     let base_offset = i32::from(self.offsets.ptr.vmmemory_definition_base());
                     let current_length_offset =
                         i32::from(self.offsets.ptr.vmmemory_definition_current_length());
-                    (memory, base_offset, current_length_offset, None)
+                    (memory, base_offset, current_length_offset, def_mt)
                 } else {
                     let owned_index = self.module.owned_memory_index(def_index);
                     let owned_base_offset =
@@ -2065,16 +2147,17 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
                 }
             } else {
                 let from_offset = self.offsets.vmctx_vmmemory_import_from(index);
-                let memory = func.create_global_value(ir::GlobalValueData::Load {
-                    base: vmctx,
-                    offset: Offset32::new(i32::try_from(from_offset).unwrap()),
-                    global_type: pointer_type,
-                    flags: MemFlags::trusted().with_readonly(),
-                });
+                let (memory, def_mt) = self.load_pointer_with_memtypes(
+                    func,
+                    vmctx,
+                    from_offset,
+                    true,
+                    self.pcc_vmctx_memtype,
+                );
                 let base_offset = i32::from(self.offsets.ptr.vmmemory_definition_base());
                 let current_length_offset =
                     i32::from(self.offsets.ptr.vmmemory_definition_current_length());
-                (memory, base_offset, current_length_offset, None)
+                (memory, base_offset, current_length_offset, def_mt)
             }
         };
 
@@ -2140,7 +2223,9 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
                                 );
                                 *size = std::cmp::max(*size, fields_end);
                             }
-                            _ => {}
+                            _ => {
+                                panic!("Bad memtype");
+                            }
                         }
                         // Apply a fact to the base pointer.
                         (Some(base_fact), Some(data_mt))
@@ -2199,7 +2284,9 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
                                     offset + u64::from(self.isa.pointer_type().bytes()),
                                 );
                             }
-                            _ => {}
+                            _ => {
+                                panic!("Bad memtype");
+                            }
                         }
                         // Apply a fact to the base pointer.
                         (Some(base_fact), Some(data_mt))
@@ -2829,16 +2916,6 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         if self.tunables.consume_fuel && state.reachable() {
             self.fuel_function_exit(builder);
         }
-        if let Some(pcc_vmctx_memtype) = self.pcc_vmctx_memtype {
-            // Sort the fields by offset in the struct definition for
-            // vmctx, now that we've completed it.
-            match &mut builder.func.memory_types[pcc_vmctx_memtype] {
-                ir::MemoryTypeData::Struct { fields, .. } => {
-                    fields.sort_by_key(|f| f.offset);
-                }
-                _ => {}
-            }
-        }
         Ok(())
     }
 
diff --git a/crates/fuzzing/src/generators/config.rs b/crates/fuzzing/src/generators/config.rs
index c974af2d7..d54202112 100644
--- a/crates/fuzzing/src/generators/config.rs
+++ b/crates/fuzzing/src/generators/config.rs
@@ -342,7 +342,6 @@ impl Config {
     pub fn disable_unimplemented_winch_proposals(&mut self) {
         self.module_config.config.simd_enabled = false;
         self.module_config.config.relaxed_simd_enabled = false;
-        self.module_config.config.memory64_enabled = false;
         self.module_config.config.gc_enabled = false;
         self.module_config.config.threads_enabled = false;
         self.module_config.config.tail_call_enabled = false;
diff --git a/crates/runtime/src/traphandlers.rs b/crates/runtime/src/traphandlers.rs
index 06855d307..92003e8b3 100644
--- a/crates/runtime/src/traphandlers.rs
+++ b/crates/runtime/src/traphandlers.rs
@@ -521,15 +521,6 @@ impl CallThreadState {
     }
 }
 
-struct ResetCell<'a, T: Copy>(&'a Cell<T>, T);
-
-impl<T: Copy> Drop for ResetCell<'_, T> {
-    #[inline]
-    fn drop(&mut self) {
-        self.0.set(self.1);
-    }
-}
-
 // A private inner module for managing the TLS state that we require across
 // calls in wasm. The WebAssembly code is called from C++ and then a trap may
 // happen which requires us to read some contextual state to figure out what to
diff --git a/crates/wasi/src/ctx.rs b/crates/wasi/src/ctx.rs
index 4e9f3578e..98233fa45 100644
--- a/crates/wasi/src/ctx.rs
+++ b/crates/wasi/src/ctx.rs
@@ -30,6 +30,7 @@ pub struct WasiCtxBuilder {
     wall_clock: Box<dyn HostWallClock + Send>,
     monotonic_clock: Box<dyn HostMonotonicClock + Send>,
     allowed_network_uses: AllowedNetworkUses,
+    allow_blocking_current_thread: bool,
     built: bool,
 }
 
@@ -80,6 +81,7 @@ impl WasiCtxBuilder {
             wall_clock: wall_clock(),
             monotonic_clock: monotonic_clock(),
             allowed_network_uses: AllowedNetworkUses::default(),
+            allow_blocking_current_thread: false,
             built: false,
         }
     }
@@ -115,6 +117,37 @@ impl WasiCtxBuilder {
         self.inherit_stdin().inherit_stdout().inherit_stderr()
     }
 
+    /// Configures whether or not blocking operations made through this
+    /// `WasiCtx` are allowed to block the current thread.
+    ///
+    /// WASI is currently implemented on top of the Rust
+    /// [Tokio](https://tokio.rs/) library. While most WASI APIs are
+    /// non-blocking some are instead blocking from the perspective of
+    /// WebAssembly. For example opening a file is a blocking operation with
+    /// respect to WebAssembly but it's implemented as an asynchronous operation
+    /// on the host. This is currently done with Tokio's
+    /// [`spawn_blocking`](https://docs.rs/tokio/latest/tokio/task/fn.spawn_blocking.html).
+    ///
+    /// When WebAssembly is used in a synchronous context, for example when
+    /// [`Config::async_support`] is disabled, then this asynchronous operation
+    /// is quickly turned back into a synchronous operation with a `block_on` in
+    /// Rust. This switching back-and-forth between a blocking a non-blocking
+    /// context can have overhead, and this option exists to help alleviate this
+    /// overhead.
+    ///
+    /// This option indicates that for WASI functions that are blocking from the
+    /// perspective of WebAssembly it's ok to block the native thread as well.
+    /// This means that this back-and-forth between async and sync won't happen
+    /// and instead blocking operations are performed on-thread (such as opening
+    /// a file). This can improve the performance of WASI operations when async
+    /// support is disabled.
+    ///
+    /// [`Config::async_support`]: https://docs.rs/wasmtime/latest/wasmtime/struct.Config.html#method.async_support
+    pub fn allow_blocking_current_thread(&mut self, enable: bool) -> &mut Self {
+        self.allow_blocking_current_thread = enable;
+        self
+    }
+
     pub fn envs(&mut self, env: &[(impl AsRef<str>, impl AsRef<str>)]) -> &mut Self {
         self.env.extend(
             env.iter()
@@ -162,7 +195,13 @@ impl WasiCtxBuilder {
             open_mode |= OpenMode::WRITE;
         }
         self.preopens.push((
-            Dir::new(dir, perms, file_perms, open_mode),
+            Dir::new(
+                dir,
+                perms,
+                file_perms,
+                open_mode,
+                self.allow_blocking_current_thread,
+            ),
             path.as_ref().to_owned(),
         ));
         self
@@ -263,6 +302,7 @@ impl WasiCtxBuilder {
             wall_clock,
             monotonic_clock,
             allowed_network_uses,
+            allow_blocking_current_thread,
             built: _,
         } = mem::replace(self, Self::new());
         self.built = true;
@@ -281,6 +321,7 @@ impl WasiCtxBuilder {
             wall_clock,
             monotonic_clock,
             allowed_network_uses,
+            allow_blocking_current_thread,
         }
     }
 
@@ -310,6 +351,7 @@ pub struct WasiCtx {
     pub(crate) stderr: Box<dyn StdoutStream>,
     pub(crate) socket_addr_check: SocketAddrCheck,
     pub(crate) allowed_network_uses: AllowedNetworkUses,
+    pub(crate) allow_blocking_current_thread: bool,
 }
 
 pub struct AllowedNetworkUses {
diff --git a/crates/wasi/src/filesystem.rs b/crates/wasi/src/filesystem.rs
index 86bd9398d..3dffa4603 100644
--- a/crates/wasi/src/filesystem.rs
+++ b/crates/wasi/src/filesystem.rs
@@ -74,6 +74,7 @@ bitflags::bitflags! {
     }
 }
 
+#[derive(Clone)]
 pub struct File {
     /// The operating system File this struct is mediating access to.
     ///
@@ -92,14 +93,22 @@ pub struct File {
     /// doesn't presently provide a cross-platform equivelant of reading the
     /// oflags back out using fcntl.
     pub open_mode: OpenMode,
+
+    allow_blocking_current_thread: bool,
 }
 
 impl File {
-    pub fn new(file: cap_std::fs::File, perms: FilePerms, open_mode: OpenMode) -> Self {
+    pub fn new(
+        file: cap_std::fs::File,
+        perms: FilePerms,
+        open_mode: OpenMode,
+        allow_blocking_current_thread: bool,
+    ) -> Self {
         Self {
             file: Arc::new(file),
             perms,
             open_mode,
+            allow_blocking_current_thread,
         }
     }
 
@@ -110,11 +119,31 @@ impl File {
         F: FnOnce(&cap_std::fs::File) -> R + Send + 'static,
         R: Send + 'static,
     {
-        let f = self.file.clone();
-        spawn_blocking(move || body(&f)).await
+        match self._spawn_blocking(body) {
+            SpawnBlocking::Done(result) => result,
+            SpawnBlocking::Spawned(task) => task.await,
+        }
+    }
+
+    fn _spawn_blocking<F, R>(&self, body: F) -> SpawnBlocking<R>
+    where
+        F: FnOnce(&cap_std::fs::File) -> R + Send + 'static,
+        R: Send + 'static,
+    {
+        if self.allow_blocking_current_thread {
+            SpawnBlocking::Done(body(&self.file))
+        } else {
+            let f = self.file.clone();
+            SpawnBlocking::Spawned(spawn_blocking(move || body(&f)))
+        }
     }
 }
 
+enum SpawnBlocking<T> {
+    Done(T),
+    Spawned(AbortOnDropJoinHandle<T>),
+}
+
 bitflags::bitflags! {
     #[derive(Copy, Clone, Debug, PartialEq, Eq)]
     pub struct DirPerms: usize {
@@ -146,6 +175,8 @@ pub struct Dir {
     /// doesn't presently provide a cross-platform equivelant of reading the
     /// oflags back out using fcntl.
     pub open_mode: OpenMode,
+
+    allow_blocking_current_thread: bool,
 }
 
 impl Dir {
@@ -154,12 +185,14 @@ impl Dir {
         perms: DirPerms,
         file_perms: FilePerms,
         open_mode: OpenMode,
+        allow_blocking_current_thread: bool,
     ) -> Self {
         Dir {
             dir: Arc::new(dir),
             perms,
             file_perms,
             open_mode,
+            allow_blocking_current_thread,
         }
     }
 
@@ -170,30 +203,38 @@ impl Dir {
         F: FnOnce(&cap_std::fs::Dir) -> R + Send + 'static,
         R: Send + 'static,
     {
-        let d = self.dir.clone();
-        spawn_blocking(move || body(&d)).await
+        if self.allow_blocking_current_thread {
+            body(&self.dir)
+        } else {
+            let d = self.dir.clone();
+            spawn_blocking(move || body(&d)).await
+        }
     }
 }
 
 pub struct FileInputStream {
-    file: Arc<cap_std::fs::File>,
+    file: File,
     position: u64,
 }
 impl FileInputStream {
-    pub fn new(file: Arc<cap_std::fs::File>, position: u64) -> Self {
-        Self { file, position }
+    pub fn new(file: &File, position: u64) -> Self {
+        Self {
+            file: file.clone(),
+            position,
+        }
     }
 
     pub async fn read(&mut self, size: usize) -> Result<Bytes, StreamError> {
         use system_interface::fs::FileIoExt;
-        let f = Arc::clone(&self.file);
         let p = self.position;
-        let (r, mut buf) = spawn_blocking(move || {
-            let mut buf = BytesMut::zeroed(size);
-            let r = f.read_at(&mut buf, p);
-            (r, buf)
-        })
-        .await;
+        let (r, mut buf) = self
+            .file
+            .spawn_blocking(move |f| {
+                let mut buf = BytesMut::zeroed(size);
+                let r = f.read_at(&mut buf, p);
+                (r, buf)
+            })
+            .await;
         let n = read_result(r)?;
         buf.truncate(n);
         self.position += n as u64;
@@ -222,7 +263,7 @@ pub(crate) enum FileOutputMode {
 }
 
 pub(crate) struct FileOutputStream {
-    file: Arc<cap_std::fs::File>,
+    file: File,
     mode: FileOutputMode,
     state: OutputState,
 }
@@ -238,16 +279,17 @@ enum OutputState {
 }
 
 impl FileOutputStream {
-    pub fn write_at(file: Arc<cap_std::fs::File>, position: u64) -> Self {
+    pub fn write_at(file: &File, position: u64) -> Self {
         Self {
-            file,
+            file: file.clone(),
             mode: FileOutputMode::Position(position),
             state: OutputState::Ready,
         }
     }
-    pub fn append(file: Arc<cap_std::fs::File>) -> Self {
+
+    pub fn append(file: &File) -> Self {
         Self {
-            file,
+            file: file.clone(),
             mode: FileOutputMode::Append,
             state: OutputState::Ready,
         }
@@ -275,33 +317,43 @@ impl HostOutputStream for FileOutputStream {
             return Ok(());
         }
 
-        let f = Arc::clone(&self.file);
         let m = self.mode;
-        let task = spawn_blocking(move || match m {
-            FileOutputMode::Position(mut p) => {
-                let mut total = 0;
-                let mut buf = buf;
-                while !buf.is_empty() {
-                    let nwritten = f.write_at(buf.as_ref(), p)?;
-                    // afterwards buf contains [nwritten, len):
-                    let _ = buf.split_to(nwritten);
-                    p += nwritten as u64;
-                    total += nwritten;
+        let result = self.file._spawn_blocking(move |f| {
+            match m {
+                FileOutputMode::Position(mut p) => {
+                    let mut total = 0;
+                    let mut buf = buf;
+                    while !buf.is_empty() {
+                        let nwritten = f.write_at(buf.as_ref(), p)?;
+                        // afterwards buf contains [nwritten, len):
+                        let _ = buf.split_to(nwritten);
+                        p += nwritten as u64;
+                        total += nwritten;
+                    }
+                    Ok(total)
                 }
-                Ok(total)
-            }
-            FileOutputMode::Append => {
-                let mut total = 0;
-                let mut buf = buf;
-                while !buf.is_empty() {
-                    let nwritten = f.append(buf.as_ref())?;
-                    let _ = buf.split_to(nwritten);
-                    total += nwritten;
+                FileOutputMode::Append => {
+                    let mut total = 0;
+                    let mut buf = buf;
+                    while !buf.is_empty() {
+                        let nwritten = f.append(buf.as_ref())?;
+                        let _ = buf.split_to(nwritten);
+                        total += nwritten;
+                    }
+                    Ok(total)
                 }
-                Ok(total)
             }
         });
-        self.state = OutputState::Waiting(task);
+        self.state = match result {
+            SpawnBlocking::Done(Ok(nwritten)) => {
+                if let FileOutputMode::Position(ref mut p) = &mut self.mode {
+                    *p += nwritten as u64;
+                }
+                OutputState::Ready
+            }
+            SpawnBlocking::Done(Err(e)) => OutputState::Error(e),
+            SpawnBlocking::Spawned(task) => OutputState::Waiting(task),
+        };
         Ok(())
     }
     fn flush(&mut self) -> Result<(), StreamError> {
diff --git a/crates/wasi/src/host/filesystem.rs b/crates/wasi/src/host/filesystem.rs
index 6094bddb7..5ad40071d 100644
--- a/crates/wasi/src/host/filesystem.rs
+++ b/crates/wasi/src/host/filesystem.rs
@@ -493,6 +493,7 @@ impl<T: WasiView> HostDescriptor for T {
         use system_interface::fs::{FdFlags, GetSetFdFlags};
         use types::{DescriptorFlags, OpenFlags};
 
+        let allow_blocking_current_thread = self.ctx().allow_blocking_current_thread;
         let table = self.table();
         let d = table.get(&fd)?.dir()?;
         if !d.perms.contains(DirPerms::READ) {
@@ -609,11 +610,15 @@ impl<T: WasiView> HostDescriptor for T {
                 d.perms,
                 d.file_perms,
                 open_mode,
+                allow_blocking_current_thread,
             )))?),
 
-            OpenResult::File(file) => {
-                Ok(table.push(Descriptor::File(File::new(file, d.file_perms, open_mode)))?)
-            }
+            OpenResult::File(file) => Ok(table.push(Descriptor::File(File::new(
+                file,
+                d.file_perms,
+                open_mode,
+                allow_blocking_current_thread,
+            )))?),
 
             OpenResult::NotDir => Err(ErrorCode::NotDirectory.into()),
         }
@@ -730,11 +735,9 @@ impl<T: WasiView> HostDescriptor for T {
         if !f.perms.contains(FilePerms::READ) {
             Err(types::ErrorCode::BadDescriptor)?;
         }
-        // Duplicate the file descriptor so that we get an indepenent lifetime.
-        let clone = std::sync::Arc::clone(&f.file);
 
         // Create a stream view for it.
-        let reader = FileInputStream::new(clone, offset);
+        let reader = FileInputStream::new(f, offset);
 
         // Insert the stream view into the table. Trap if the table is full.
         let index = self.table().push(InputStream::File(reader))?;
@@ -754,11 +757,8 @@ impl<T: WasiView> HostDescriptor for T {
             Err(types::ErrorCode::BadDescriptor)?;
         }
 
-        // Duplicate the file descriptor so that we get an indepenent lifetime.
-        let clone = std::sync::Arc::clone(&f.file);
-
         // Create a stream view for it.
-        let writer = FileOutputStream::write_at(clone, offset);
+        let writer = FileOutputStream::write_at(f, offset);
         let writer: OutputStream = Box::new(writer);
 
         // Insert the stream view into the table. Trap if the table is full.
@@ -777,11 +777,9 @@ impl<T: WasiView> HostDescriptor for T {
         if !f.perms.contains(FilePerms::WRITE) {
             Err(types::ErrorCode::BadDescriptor)?;
         }
-        // Duplicate the file descriptor so that we get an indepenent lifetime.
-        let clone = std::sync::Arc::clone(&f.file);
 
         // Create a stream view for it.
-        let appender = FileOutputStream::append(clone);
+        let appender = FileOutputStream::append(f);
         let appender: OutputStream = Box::new(appender);
 
         // Insert the stream view into the table. Trap if the table is full.
diff --git a/crates/wasmtime/src/config.rs b/crates/wasmtime/src/config.rs
index bef3e7061..271c9bc25 100644
--- a/crates/wasmtime/src/config.rs
+++ b/crates/wasmtime/src/config.rs
@@ -1009,9 +1009,14 @@ impl Config {
     /// Configures whether Cranelift should perform a NaN-canonicalization pass.
     ///
     /// When Cranelift is used as a code generation backend this will configure
-    /// it to replace NaNs with a single canonical value. This is useful for users
-    /// requiring entirely deterministic WebAssembly computation.
-    /// This is not required by the WebAssembly spec, so it is not enabled by default.
+    /// it to replace NaNs with a single canonical value. This is useful for
+    /// users requiring entirely deterministic WebAssembly computation.  This is
+    /// not required by the WebAssembly spec, so it is not enabled by default.
+    ///
+    /// Note that this option affects not only WebAssembly's `f32` and `f64`
+    /// types but additionally the `v128` type. This option will cause
+    /// operations using any of these types to have extra checks placed after
+    /// them to normalize NaN values as needed.
     ///
     /// The default value for this is `false`
     #[cfg(any(feature = "cranelift", feature = "winch"))]
diff --git a/src/commands/run.rs b/src/commands/run.rs
index 1413874ee..2c1eed303 100644
--- a/src/commands/run.rs
+++ b/src/commands/run.rs
@@ -207,15 +207,26 @@ impl RunCommand {
             }
         }
 
+        // Pre-emptively initialize and install a Tokio runtime ambiently in the
+        // environment when executing the module. Without this whenever a WASI
+        // call is made that needs to block on a future a Tokio runtime is
+        // configured and entered, and this appears to be slower than simply
+        // picking an existing runtime out of the environment and using that.
+        // The goal of this is to improve the performance of WASI-related
+        // operations that block in the CLI since the CLI doesn't use async to
+        // invoke WebAssembly.
+        let result = wasmtime_wasi::runtime::with_ambient_tokio_runtime(|| {
+            self.load_main_module(&mut store, &mut linker, &main, modules)
+                .with_context(|| {
+                    format!(
+                        "failed to run main module `{}`",
+                        self.module_and_args[0].to_string_lossy()
+                    )
+                })
+        });
+
         // Load the main wasm module.
-        match self
-            .load_main_module(&mut store, &mut linker, &main, modules)
-            .with_context(|| {
-                format!(
-                    "failed to run main module `{}`",
-                    self.module_and_args[0].to_string_lossy()
-                )
-            }) {
+        match result {
             Ok(()) => (),
             Err(e) => {
                 // Exit the process if Wasmtime understands the error;
@@ -799,6 +810,12 @@ impl RunCommand {
         let mut builder = wasmtime_wasi::WasiCtxBuilder::new();
         builder.inherit_stdio().args(&self.compute_argv()?);
 
+        // It's ok to block the current thread since we're the only thread in
+        // the program as the CLI. This helps improve the performance of some
+        // blocking operations in WASI, for example, by skipping the
+        // back-and-forth between sync and async.
+        builder.allow_blocking_current_thread(true);
+
         if self.run.common.wasi.inherit_env == Some(true) {
             for (k, v) in std::env::vars() {
                 builder.env(&k, &v);
diff --git a/winch/codegen/src/codegen/env.rs b/winch/codegen/src/codegen/env.rs
index 978d16561..53e82fa46 100644
--- a/winch/codegen/src/codegen/env.rs
+++ b/winch/codegen/src/codegen/env.rs
@@ -328,8 +328,7 @@ impl<'a, 'translation, 'data, P: PtrSize> FuncEnv<'a, 'translation, 'data, P> {
                     current_length_offset,
                     style,
                     ty: if plan.memory.memory64 {
-                        // TODO: Add support for 64-bit memories.
-                        unimplemented!("memory64")
+                        WasmValType::I64
                     } else {
                         WasmValType::I32
                     },
diff --git a/winch/codegen/src/visitor.rs b/winch/codegen/src/visitor.rs
index 3ad1ec585..da9051302 100644
--- a/winch/codegen/src/visitor.rs
+++ b/winch/codegen/src/visitor.rs
@@ -10,6 +10,7 @@ use crate::masm::{
     DivKind, ExtendKind, FloatCmpKind, IntCmpKind, MacroAssembler, MemMoveDirection, OperandSize,
     RegImm, RemKind, RoundingMode, SPOffset, ShiftKind, TruncKind,
 };
+use crate::reg::Reg;
 use crate::stack::{TypedReg, Val};
 use cranelift_codegen::ir::TrapCode;
 use regalloc2::RegClass;
@@ -1593,8 +1594,24 @@ where
         //   [ vmctx, delta, index ]
         self.context.stack.extend([mem.try_into().unwrap()]);
 
+        let heap = self.env.resolve_heap(MemoryIndex::from_u32(mem));
         let builtin = self.env.builtins.memory32_grow::<M::ABI, M::Ptr>();
-        FnCall::emit::<M, M::Ptr>(self.masm, &mut self.context, Callee::Builtin(builtin))
+        FnCall::emit::<M, M::Ptr>(self.masm, &mut self.context, Callee::Builtin(builtin));
+
+        // The memory32_grow builtin returns a pointer type, therefore we must
+        // ensure that the return type is representative of the address space of
+        // the heap type.
+        match (self.env.ptr_type(), heap.ty) {
+            (WasmValType::I64, WasmValType::I64) => {}
+            // When the heap type is smaller than the pointer type, we adjust
+            // the result of the memory32_grow builtin.
+            (WasmValType::I64, WasmValType::I32) => {
+                let top: Reg = self.context.pop_to_reg(self.masm, None).into();
+                self.masm.wrap(top.into(), top.into());
+                self.context.stack.push(TypedReg::i32(top).into());
+            }
+            _ => unimplemented!("Support for 32-bit platforms"),
+        }
     }
 
     fn visit_data_drop(&mut self, data_index: u32) {


Test patch:
diff --git a/tests/disas/pcc-imported-memory.wat b/tests/disas/pcc-imported-memory.wat
new file mode 100644
index 000000000..5bd75ad01
--- /dev/null
+++ b/tests/disas/pcc-imported-memory.wat
@@ -0,0 +1,78 @@
+;;! target = "x86_64"
+;;! test = "compile"
+;;! flags = [ "-Oopt-level=2", "-Cpcc=y" ]
+
+(module
+  (type (;0;) (func))
+  (import "" "" (memory (;0;) 1))
+  (func (;0;) (type 0)
+    (local i32 i32)
+    memory.size
+    local.set 0
+    block  ;; label = @1
+      block  ;; label = @2
+        memory.size
+        i32.const 65536
+        i32.mul
+        i32.const 4
+        local.get 0
+        i32.add
+        i32.le_u
+        br_if 0 (;@2;)
+        local.get 0
+        i32.const 0
+        i32.le_s
+        br_if 0 (;@2;)
+        local.get 0
+        i32.load align=1
+        local.set 1
+        br 1 (;@1;)
+      end
+      i32.const 0
+      local.set 1
+    end
+    local.get 1
+    drop))
+
+;; function u0:0:
+;;   pushq   %rbp
+;;   unwind PushFrameRegs { offset_upward_to_caller_sp: 16 }
+;;   movq    %rsp, %rbp
+;;   movq    8(%rdi), %r10
+;;   movq    0(%r10), %r10
+;;   cmpq    %rsp, %r10
+;;   jnbe #trap=stk_ovf
+;;   unwind DefineNewFrame { offset_upward_to_caller_sp: 16, offset_downward_to_clobbers: 0 }
+;; block0:
+;;   movq    72(%rdi), %r10
+;;   movq    8(%r10), %rax
+;;   movl    $65536, %r10d
+;;   xorq    %rdx, %rdx, %rdx
+;;   div     %rax, %rdx, %r10, %rax, %rdx ; trap=int_divz
+;;   movq    %rax, %r9
+;;   shll    $16, %r9d, %r9d
+;;   lea     4(%rax), %r10d
+;;   cmpl    %r10d, %r9d
+;;   jbe     label1; j label2
+;; block1:
+;;   jmp     label5
+;; block2:
+;;   testl   %eax, %eax
+;;   jle     label3; j label4
+;; block3:
+;;   jmp     label5
+;; block4:
+;;   movq    72(%rdi), %rcx
+;;   movq    0(%rcx), %rcx
+;;   movl    %eax, %edx
+;;   movl    0(%rcx,%rdx,1), %edi
+;;   jmp     label6
+;; block5:
+;;   xorl    %edi, %edi, %edi
+;;   jmp     label6
+;; block6:
+;;   jmp     label7
+;; block7:
+;;   movq    %rbp, %rsp
+;;   popq    %rbp
+;;   ret
diff --git a/winch/filetests/filetests/x64/load/grow_load.wat b/winch/filetests/filetests/x64/load/grow_load.wat
new file mode 100644
index 000000000..1e1bf5e51
--- /dev/null
+++ b/winch/filetests/filetests/x64/load/grow_load.wat
@@ -0,0 +1,103 @@
+;;! target = "x86_64"
+(module
+  (type (;0;) (func (param f32 f64 i64 i64 f64 f64 f32) (result f32 f64 f32)))
+  (func (;0;) (type 0) (param f32 f64 i64 i64 f64 f64 f32) (result f32 f64 f32)
+    global.get 1
+    i32.eqz
+    if ;; label = @1
+      unreachable
+    end
+    global.get 1
+    i32.const 1
+    i32.sub
+    global.set 1
+    memory.size
+    memory.grow
+    i64.load8_s offset=143396
+    (drop)
+    (f32.const 0)
+    (f64.const 0)
+    (f32.const 0)
+  )
+  (memory (;1;) 10 10)
+  (global (;0;) f32 f32.const 0x1.d6a0d6p+87 (;=284477330000000000000000000;))
+  (global (;1;) (mut i32) i32.const 1000)
+  (export "main" (func 0))
+  (export "0" (memory 0))
+  (export "1" (global 0))
+)
+;;      	 55                   	push	rbp
+;;      	 4889e5               	mov	rbp, rsp
+;;      	 4c8b5f08             	mov	r11, qword ptr [rdi + 8]
+;;      	 4d8b1b               	mov	r11, qword ptr [r11]
+;;      	 4981c360000000       	add	r11, 0x60
+;;      	 4939e3               	cmp	r11, rsp
+;;      	 0f87f1000000         	ja	0x10c
+;;   1b:	 4989fe               	mov	r14, rdi
+;;      	 4883ec50             	sub	rsp, 0x50
+;;      	 48897c2448           	mov	qword ptr [rsp + 0x48], rdi
+;;      	 4889742440           	mov	qword ptr [rsp + 0x40], rsi
+;;      	 f30f1144243c         	movss	dword ptr [rsp + 0x3c], xmm0
+;;      	 f20f114c2430         	movsd	qword ptr [rsp + 0x30], xmm1
+;;      	 4889542428           	mov	qword ptr [rsp + 0x28], rdx
+;;      	 48894c2420           	mov	qword ptr [rsp + 0x20], rcx
+;;      	 f20f11542418         	movsd	qword ptr [rsp + 0x18], xmm2
+;;      	 f20f115c2410         	movsd	qword ptr [rsp + 0x10], xmm3
+;;      	 f30f1164240c         	movss	dword ptr [rsp + 0xc], xmm4
+;;      	 4c890424             	mov	qword ptr [rsp], r8
+;;      	 418b4670             	mov	eax, dword ptr [r14 + 0x70]
+;;      	 83f800               	cmp	eax, 0
+;;      	 b800000000           	mov	eax, 0
+;;      	 400f94c0             	sete	al
+;;      	 85c0                 	test	eax, eax
+;;      	 0f8402000000         	je	0x72
+;;   70:	 0f0b                 	ud2	
+;;      	 418b4670             	mov	eax, dword ptr [r14 + 0x70]
+;;      	 83e801               	sub	eax, 1
+;;      	 41894670             	mov	dword ptr [r14 + 0x70], eax
+;;      	 498b4658             	mov	rax, qword ptr [r14 + 0x58]
+;;      	 c1e810               	shr	eax, 0x10
+;;      	 4883ec04             	sub	rsp, 4
+;;      	 890424               	mov	dword ptr [rsp], eax
+;;      	 4d8b5e38             	mov	r11, qword ptr [r14 + 0x38]
+;;      	 498b0b               	mov	rcx, qword ptr [r11]
+;;      	 4883ec0c             	sub	rsp, 0xc
+;;      	 4c89f7               	mov	rdi, r14
+;;      	 8b74240c             	mov	esi, dword ptr [rsp + 0xc]
+;;      	 ba00000000           	mov	edx, 0
+;;      	 ffd1                 	call	rcx
+;;      	 4883c40c             	add	rsp, 0xc
+;;      	 4883c404             	add	rsp, 4
+;;      	 4c8b742448           	mov	r14, qword ptr [rsp + 0x48]
+;;      	 89c0                 	mov	eax, eax
+;;      	 498b4e50             	mov	rcx, qword ptr [r14 + 0x50]
+;;      	 4801c1               	add	rcx, rax
+;;      	 4881c124300200       	add	rcx, 0x23024
+;;      	 480fbe01             	movsx	rax, byte ptr [rcx]
+;;      	 f30f100543000000     	movss	xmm0, dword ptr [rip + 0x43]
+;;      	 4883ec0c             	sub	rsp, 0xc
+;;      	 f2440f103d3e000000   	
+;; 				movsd	xmm15, qword ptr [rip + 0x3e]
+;;      	 f2440f113c24         	movsd	qword ptr [rsp], xmm15
+;;      	 f3440f103d27000000   	
+;; 				movss	xmm15, dword ptr [rip + 0x27]
+;;      	 f3440f117c2408       	movss	dword ptr [rsp + 8], xmm15
+;;      	 488b44240c           	mov	rax, qword ptr [rsp + 0xc]
+;;      	 415b                 	pop	r11
+;;      	 4c8918               	mov	qword ptr [rax], r11
+;;      	 448b1c24             	mov	r11d, dword ptr [rsp]
+;;      	 4883c404             	add	rsp, 4
+;;      	 44895808             	mov	dword ptr [rax + 8], r11d
+;;      	 4883c450             	add	rsp, 0x50
+;;      	 5d                   	pop	rbp
+;;      	 c3                   	ret	
+;;  10c:	 0f0b                 	ud2	
+;;  10e:	 0000                 	add	byte ptr [rax], al
+;;  110:	 0000                 	add	byte ptr [rax], al
+;;  112:	 0000                 	add	byte ptr [rax], al
+;;  114:	 0000                 	add	byte ptr [rax], al
+;;  116:	 0000                 	add	byte ptr [rax], al
+;;  118:	 0000                 	add	byte ptr [rax], al
+;;  11a:	 0000                 	add	byte ptr [rax], al
+;;  11c:	 0000                 	add	byte ptr [rax], al
+;;  11e:	 0000                 	add	byte ptr [rax], al