Cooperative Multithreading (#11751)

* Initial work

* Almost compiling

* Partially working

* Cleanup

* Fix merge

* Cancellation and suspension refactoring

* Remove printlns

* Test with several threads

* More testing

* Cancellation

* Fix cancellation for explicit suspends

* Finish cancellation test

* Store threads in the instance table

* Deletion almost there

* Tests all pass

* Tighten up task deletion

* Set thread state correctly

* Store pairs of thread and task ids

* Remove lift abi members

* Cleanup unnecessary change

* More cleanup

* Cleanup

* Revert cargo changes

* Revert cargo changes

* Comments

* Comments on test

* Update comments

* Update comments

* Add space that was removed in an earlier commit

* Cleanup

* Delete threads from the instance table

* Revert cargo file

* Remove dead code

* Clippy changes

* Clippy

* Revert unnecessary changes

* Revert unnecessary changes

* Revert unnecessary changes

* Revert unnecessary changes

* Revert unnecessary changes

* Review comments

* Review feedback

* Make thread IDs per-component-instance

* Fix config

* Tighten up completion

* Clippy

* Trigger full PR test

* Move funcref table reading

* Remove unused import

* Formatting

* Rename RemoveOnDrop

* Review feedback

* Review feedback

* Move start thread closure

* Review feedback

* Review feedback

* Correct feature for import

* Review feedback

* Disable failing tests

* Enable fixed tests

* Review feedback

* Readd tests

* Ignore task deletion test with Miri

Author: TartanLlama

crates/cli-flags/src/lib.rs

@@ -383,6 +383,9 @@ wasmtime_option_group! {
         /// Component model support for async lifting/lowering: this corresponds
         /// to the 🚟 emoji in the component model specification.
         pub component_model_async_stackful: Option<bool>,
+        /// Component model support for threading: this corresponds
+        /// to the 🧵 emoji in the component model specification.
+        pub component_model_threading: Option<bool>,
         /// Component model support for `error-context`: this corresponds
         /// to the 📝 emoji in the component model specification.
         pub component_model_error_context: Option<bool>,
@@ -1059,6 +1062,7 @@ impl CommonOptions {
             ("component-model-async", component_model_async, wasm_component_model_async)
             ("component-model-async", component_model_async_builtins, wasm_component_model_async_builtins)
             ("component-model-async", component_model_async_stackful, wasm_component_model_async_stackful)
+            ("component-model-async", component_model_threading, wasm_component_model_threading)
             ("component-model", component_model_error_context, wasm_component_model_error_context)
             ("threads", threads, wasm_threads)
             ("gc", gc, wasm_gc)

crates/cranelift/src/compiler/component.rs

@@ -871,6 +871,94 @@ impl<'a> TrampolineCompiler<'a> {
                     },
                 );
             }
+            Trampoline::ThreadIndex => {
+                self.translate_libcall(
+                    host::thread_index,
+                    TrapSentinel::NegativeOne,
+                    WasmArgs::InRegisters,
+                    |_, _| {},
+                );
+            }
+            Trampoline::ThreadNewIndirect {
+                instance,
+                start_func_table_idx,
+                start_func_ty_idx,
+            } => {
+                self.translate_libcall(
+                    host::thread_new_indirect,
+                    TrapSentinel::NegativeOne,
+                    WasmArgs::InRegisters,
+                    |me, params| {
+                        params.push(me.index_value(*instance));
+                        params.push(me.index_value(*start_func_table_idx));
+                        params.push(me.index_value(*start_func_ty_idx));
+                    },
+                );
+            }
+            Trampoline::ThreadSwitchTo {
+                instance,
+                cancellable,
+            } => {
+                self.translate_libcall(
+                    host::thread_switch_to,
+                    TrapSentinel::NegativeOne,
+                    WasmArgs::InRegisters,
+                    |me, params| {
+                        params.push(me.index_value(*instance));
+                        params.push(
+                            me.builder
+                                .ins()
+                                .iconst(ir::types::I8, i64::from(*cancellable)),
+                        );
+                    },
+                );
+            }
+            Trampoline::ThreadSuspend {
+                instance,
+                cancellable,
+            } => {
+                self.translate_libcall(
+                    host::thread_suspend,
+                    TrapSentinel::NegativeOne,
+                    WasmArgs::InRegisters,
+                    |me, params| {
+                        params.push(me.index_value(*instance));
+                        params.push(
+                            me.builder
+                                .ins()
+                                .iconst(ir::types::I8, i64::from(*cancellable)),
+                        );
+                    },
+                );
+            }
+            Trampoline::ThreadResumeLater { instance } => {
+                self.translate_libcall(
+                    host::thread_resume_later,
+                    TrapSentinel::Falsy,
+                    WasmArgs::InRegisters,
+                    |me, params| {
+                        params.push(me.index_value(*instance));
+                    },
+                );
+            }
+            Trampoline::ThreadYieldTo {
+                instance,
+                cancellable,
+            } => {
+                self.translate_libcall(
+                    host::thread_yield_to,
+                    TrapSentinel::NegativeOne,
+                    WasmArgs::InRegisters,
+                    |me, params| {
+                        params.push(me.index_value(*instance));
+                        params.push(
+                            me.builder
+                                .ins()
+                                .iconst(ir::types::I8, i64::from(*cancellable)),
+                        );
+                    },
+                );
+            }
         }
     }
 

crates/environ/src/component.rs

@@ -189,6 +189,18 @@ macro_rules! foreach_builtin_component_function {
             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;
+            #[cfg(feature = "component-model-async")]
+            thread_switch_to(vmctx: vmctx, caller_instance: u32, cancellable: u8, thread_idx: u32) -> u32;
+            #[cfg(feature = "component-model-async")]
+            thread_suspend(vmctx: vmctx, caller_instance: u32, cancellable: u8) -> u32;
+            #[cfg(feature = "component-model-async")]
+            thread_resume_later(vmctx: vmctx, caller_instance: u32, thread_idx: u32) -> bool;
+            #[cfg(feature = "component-model-async")]
+            thread_yield_to(vmctx: vmctx, caller_instance: u32, cancellable: u8, thread_idx: u32) -> u32;
 
             trap(vmctx: vmctx, code: u8) -> bool;
 

crates/environ/src/component/dfg.rs

@@ -69,9 +69,12 @@ pub struct ComponentDfg {
     /// Same as `reallocs`, but for post-return.
     pub post_returns: Intern<PostReturnId, CoreDef>,
 
-    /// Same as `reallocs`, but for post-return.
+    /// Same as `reallocs`, but for memories.
     pub memories: Intern<MemoryId, CoreExport<MemoryIndex>>,
 
+    /// Same as `reallocs`, but for tables.
+    pub tables: Intern<TableId, CoreExport<TableIndex>>,
+
     /// Metadata about identified fused adapters.
     ///
     /// Note that this list is required to be populated in-order where the
@@ -483,6 +486,27 @@ pub enum Trampoline {
         instance: RuntimeComponentInstanceIndex,
         slot: u32,
     },
+    ThreadIndex,
+    ThreadNewIndirect {
+        instance: RuntimeComponentInstanceIndex,
+        start_func_ty_idx: ComponentTypeIndex,
+        start_func_table_id: TableId,
+    },
+    ThreadSwitchTo {
+        instance: RuntimeComponentInstanceIndex,
+        cancellable: bool,
+    },
+    ThreadSuspend {
+        instance: RuntimeComponentInstanceIndex,
+        cancellable: bool,
+    },
+    ThreadResumeLater {
+        instance: RuntimeComponentInstanceIndex,
+    },
+    ThreadYieldTo {
+        instance: RuntimeComponentInstanceIndex,
+        cancellable: bool,
+    },
 }
 
 #[derive(Copy, Clone, Hash, Eq, PartialEq)]
@@ -833,6 +857,15 @@ impl LinearizeDfg<'_> {
         )
     }
 
+    fn runtime_table(&mut self, table: TableId) -> RuntimeTableIndex {
+        self.intern(
+            table,
+            |me| &mut me.runtime_tables,
+            |me, table| me.core_export(&me.dfg.tables[table]),
+            |index, export| GlobalInitializer::ExtractTable(ExtractTable { index, export }),
+        )
+    }
+
     fn runtime_realloc(&mut self, realloc: ReallocId) -> RuntimeReallocIndex {
         self.intern(
             realloc,
@@ -1135,6 +1168,40 @@ impl LinearizeDfg<'_> {
                 instance: *instance,
                 slot: *slot,
             },
+            Trampoline::ThreadIndex => info::Trampoline::ThreadIndex,
+            Trampoline::ThreadNewIndirect {
+                instance,
+                start_func_ty_idx,
+                start_func_table_id,
+            } => info::Trampoline::ThreadNewIndirect {
+                instance: *instance,
+                start_func_ty_idx: *start_func_ty_idx,
+                start_func_table_idx: self.runtime_table(*start_func_table_id),
+            },
+            Trampoline::ThreadSwitchTo {
+                instance,
+                cancellable,
+            } => info::Trampoline::ThreadSwitchTo {
+                instance: *instance,
+                cancellable: *cancellable,
+            },
+            Trampoline::ThreadSuspend {
+                instance,
+                cancellable,
+            } => info::Trampoline::ThreadSuspend {
+                instance: *instance,
+                cancellable: *cancellable,
+            },
+            Trampoline::ThreadResumeLater { instance } => info::Trampoline::ThreadResumeLater {
+                instance: *instance,
+            },
+            Trampoline::ThreadYieldTo {
+                instance,
+                cancellable,
+            } => info::Trampoline::ThreadYieldTo {
+                instance: *instance,
+                cancellable: *cancellable,
+            },
         };
         let i1 = self.trampolines.push(*signature);
         let i2 = self.trampoline_defs.push(trampoline);

crates/environ/src/component/info.rs

@@ -831,7 +831,7 @@ pub enum Trampoline {
     ThreadYield {
         /// The specific component instance which is calling the intrinsic.
         instance: RuntimeComponentInstanceIndex,
-        /// If `true`, indicates the caller instance maybe receive notification
+        /// If `true`, indicates the caller instance may receive notification
         /// of task cancellation.
         cancellable: bool,
     },
@@ -1123,6 +1123,52 @@ pub enum Trampoline {
         /// Which slot to update.
         slot: u32,
     },
+
+    /// Intrinsic used to implement the `thread.index` component model builtin.
+    ThreadIndex,
+
+    /// Intrinsic used to implement the `thread.new_indirect` component model builtin.
+    ThreadNewIndirect {
+        /// The specific component instance which is calling the intrinsic.
+        instance: RuntimeComponentInstanceIndex,
+        /// The type index for the start function of the thread.
+        start_func_ty_idx: ComponentTypeIndex,
+        /// The index of the table that stores the start function.
+        start_func_table_idx: RuntimeTableIndex,
+    },
+
+    /// Intrinsic used to implement the `thread.switch-to` component model builtin.
+    ThreadSwitchTo {
+        /// The specific component instance which is calling the intrinsic.
+        instance: RuntimeComponentInstanceIndex,
+        /// If `true`, indicates the caller instance may receive notification
+        /// of task cancellation.
+        cancellable: bool,
+    },
+
+    /// Intrinsic used to implement the `thread.suspend` component model builtin.
+    ThreadSuspend {
+        /// The specific component instance which is calling the intrinsic.
+        instance: RuntimeComponentInstanceIndex,
+        /// If `true`, indicates the caller instance may receive notification
+        /// of task cancellation.
+        cancellable: bool,
+    },
+
+    /// Intrinsic used to implement the `thread.resume-later` component model builtin.
+    ThreadResumeLater {
+        /// The specific component instance which is calling the intrinsic.
+        instance: RuntimeComponentInstanceIndex,
+    },
+
+    /// Intrinsic used to implement the `thread.yield-to` component model builtin.
+    ThreadYieldTo {
+        /// The specific component instance which is calling the intrinsic.
+        instance: RuntimeComponentInstanceIndex,
+        /// If `true`, indicates the caller instance may receive notification
+        /// of task cancellation.
+        cancellable: bool,
+    },
 }
 
 impl Trampoline {
@@ -1188,6 +1234,12 @@ impl Trampoline {
             ErrorContextTransfer => format!("error-context-transfer"),
             ContextGet { .. } => format!("context-get"),
             ContextSet { .. } => format!("context-set"),
+            ThreadIndex => format!("thread-index"),
+            ThreadNewIndirect { .. } => format!("thread-new-indirect"),
+            ThreadSwitchTo { .. } => format!("thread-switch-to"),
+            ThreadSuspend { .. } => format!("thread-suspend"),
+            ThreadResumeLater { .. } => format!("thread-resume-later"),
+            ThreadYieldTo { .. } => format!("thread-yield-to"),
         }
     }
 }

crates/environ/src/component/translate.rs

@@ -318,6 +318,29 @@ enum LocalInitializer<'data> {
         func: ModuleInternedTypeIndex,
         i: u32,
     },
+    ThreadIndex {
+        func: ModuleInternedTypeIndex,
+    },
+    ThreadNewIndirect {
+        func: ModuleInternedTypeIndex,
+        start_func_ty: ComponentTypeIndex,
+        start_func_table_index: TableIndex,
+    },
+    ThreadSwitchTo {
+        func: ModuleInternedTypeIndex,
+        cancellable: bool,
+    },
+    ThreadSuspend {
+        func: ModuleInternedTypeIndex,
+        cancellable: bool,
+    },
+    ThreadResumeLater {
+        func: ModuleInternedTypeIndex,
+    },
+    ThreadYieldTo {
+        func: ModuleInternedTypeIndex,
+        cancellable: bool,
+    },
 
     // core wasm modules
     ModuleStatic(StaticModuleIndex, ComponentCoreModuleTypeId),
@@ -1121,24 +1144,42 @@ impl<'a, 'data> Translator<'a, 'data> {
                             core_func_index += 1;
                             LocalInitializer::ContextSet { i, func }
                         }
-
                         wasmparser::CanonicalFunction::ThreadIndex => {
-                            bail!("unimplemented `thread.index`");
+                            let func = self.core_func_signature(core_func_index)?;
+                            core_func_index += 1;
+                            LocalInitializer::ThreadIndex { func }
                         }
-                        wasmparser::CanonicalFunction::ThreadNewIndirect { .. } => {
-                            bail!("unimplemented `thread.new-indirect`");
+                        wasmparser::CanonicalFunction::ThreadNewIndirect {
+                            func_ty_index,
+                            table_index,
+                        } => {
+                            let func = self.core_func_signature(core_func_index)?;
+                            core_func_index += 1;
+                            LocalInitializer::ThreadNewIndirect {
+                                func,
+                                start_func_ty: ComponentTypeIndex::from_u32(func_ty_index),
+                                start_func_table_index: TableIndex::from_u32(table_index),
+                            }
                         }
-                        wasmparser::CanonicalFunction::ThreadSwitchTo { .. } => {
-                            bail!("unimplemented `thread.switch-to`");
+                        wasmparser::CanonicalFunction::ThreadSwitchTo { cancellable } => {
+                            let func = self.core_func_signature(core_func_index)?;
+                            core_func_index += 1;
+                            LocalInitializer::ThreadSwitchTo { func, cancellable }
                         }
-                        wasmparser::CanonicalFunction::ThreadSuspend { .. } => {
-                            bail!("unimplemented `thread.suspend`");
+                        wasmparser::CanonicalFunction::ThreadSuspend { cancellable } => {
+                            let func = self.core_func_signature(core_func_index)?;
+                            core_func_index += 1;
+                            LocalInitializer::ThreadSuspend { func, cancellable }
                         }
                         wasmparser::CanonicalFunction::ThreadResumeLater => {
-                            bail!("unimplemented `thread.resume-later`");
+                            let func = self.core_func_signature(core_func_index)?;
+                            core_func_index += 1;
+                            LocalInitializer::ThreadResumeLater { func }
                         }
-                        wasmparser::CanonicalFunction::ThreadYieldTo { .. } => {
-                            bail!("unimplemented `thread.yield-to`");
+                        wasmparser::CanonicalFunction::ThreadYieldTo { cancellable } => {
+                            let func = self.core_func_signature(core_func_index)?;
+                            core_func_index += 1;
+                            LocalInitializer::ThreadYieldTo { func, cancellable }
                         }
                     };
                     self.result.initializers.push(init);