[Cranelift][Atomics] Add address folding for atomic notify/wait. (#2556)

* fold address in wasm wait and notify ops

* add atomics addr folding tests

crates/cranelift/src/func_environ.rs

@@ -81,6 +81,10 @@ macro_rules! declare_function_signatures {
                 AbiParam::new(I32).uext()
             }
 
+            fn i64(&self) -> AbiParam {
+                AbiParam::new(I64)
+            }
+
             $(
                 fn $name(&mut self, func: &mut Function) -> ir::SigRef {
                     let sig = self.$name.unwrap_or_else(|| {
@@ -258,6 +262,70 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
         }
     }
 
+    fn get_memory_atomic_notify(
+        &mut self,
+        func: &mut Function,
+        memory_index: MemoryIndex,
+    ) -> (ir::SigRef, usize, BuiltinFunctionIndex) {
+        if let Some(defined_memory_index) = self.module.defined_memory_index(memory_index) {
+            (
+                self.builtin_function_signatures.memory_atomic_notify(func),
+                defined_memory_index.index(),
+                BuiltinFunctionIndex::memory_atomic_notify(),
+            )
+        } else {
+            (
+                self.builtin_function_signatures
+                    .imported_memory_atomic_notify(func),
+                memory_index.index(),
+                BuiltinFunctionIndex::imported_memory_atomic_notify(),
+            )
+        }
+    }
+
+    fn get_memory_atomic_wait(
+        &mut self,
+        func: &mut Function,
+        memory_index: MemoryIndex,
+        ty: ir::Type,
+    ) -> (ir::SigRef, usize, BuiltinFunctionIndex) {
+        match ty {
+            I32 => {
+                if let Some(defined_memory_index) = self.module.defined_memory_index(memory_index) {
+                    (
+                        self.builtin_function_signatures.memory_atomic_wait32(func),
+                        defined_memory_index.index(),
+                        BuiltinFunctionIndex::memory_atomic_wait32(),
+                    )
+                } else {
+                    (
+                        self.builtin_function_signatures
+                            .imported_memory_atomic_wait32(func),
+                        memory_index.index(),
+                        BuiltinFunctionIndex::imported_memory_atomic_wait32(),
+                    )
+                }
+            }
+            I64 => {
+                if let Some(defined_memory_index) = self.module.defined_memory_index(memory_index) {
+                    (
+                        self.builtin_function_signatures.memory_atomic_wait64(func),
+                        defined_memory_index.index(),
+                        BuiltinFunctionIndex::memory_atomic_wait64(),
+                    )
+                } else {
+                    (
+                        self.builtin_function_signatures
+                            .imported_memory_atomic_wait64(func),
+                        memory_index.index(),
+                        BuiltinFunctionIndex::imported_memory_atomic_wait64(),
+                    )
+                }
+            }
+            x => panic!("get_memory_atomic_wait unsupported type: {:?}", x),
+        }
+    }
+
     fn get_memory_init_func(&mut self, func: &mut Function) -> (ir::SigRef, BuiltinFunctionIndex) {
         (
             self.builtin_function_signatures.memory_init(func),
@@ -1368,29 +1436,50 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
 
     fn translate_atomic_wait(
         &mut self,
-        _pos: FuncCursor,
-        _index: MemoryIndex,
+        mut pos: FuncCursor,
+        memory_index: MemoryIndex,
         _heap: ir::Heap,
-        _addr: ir::Value,
-        _expected: ir::Value,
-        _timeout: ir::Value,
+        addr: ir::Value,
+        expected: ir::Value,
+        timeout: ir::Value,
     ) -> WasmResult<ir::Value> {
-        Err(WasmError::Unsupported(
-            "wasm atomics (fn translate_atomic_wait)".to_string(),
-        ))
+        let implied_ty = pos.func.dfg.value_type(expected);
+        let (func_sig, memory_index, func_idx) =
+            self.get_memory_atomic_wait(&mut pos.func, memory_index, implied_ty);
+
+        let memory_index_arg = pos.ins().iconst(I32, memory_index as i64);
+
+        let (vmctx, func_addr) = self.translate_load_builtin_function_address(&mut pos, func_idx);
+
+        let call_inst = pos.ins().call_indirect(
+            func_sig,
+            func_addr,
+            &[vmctx, memory_index_arg, addr, expected, timeout],
+        );
+
+        Ok(*pos.func.dfg.inst_results(call_inst).first().unwrap())
     }
 
     fn translate_atomic_notify(
         &mut self,
-        _pos: FuncCursor,
-        _index: MemoryIndex,
+        mut pos: FuncCursor,
+        memory_index: MemoryIndex,
         _heap: ir::Heap,
-        _addr: ir::Value,
-        _count: ir::Value,
+        addr: ir::Value,
+        count: ir::Value,
     ) -> WasmResult<ir::Value> {
-        Err(WasmError::Unsupported(
-            "wasm atomics (fn translate_atomic_notify)".to_string(),
-        ))
+        let (func_sig, memory_index, func_idx) =
+            self.get_memory_atomic_notify(&mut pos.func, memory_index);
+
+        let memory_index_arg = pos.ins().iconst(I32, memory_index as i64);
+
+        let (vmctx, func_addr) = self.translate_load_builtin_function_address(&mut pos, func_idx);
+
+        let call_inst =
+            pos.ins()
+                .call_indirect(func_sig, func_addr, &[vmctx, memory_index_arg, addr, count]);
+
+        Ok(*pos.func.dfg.inst_results(call_inst).first().unwrap())
     }
 
     fn translate_loop_header(&mut self, mut pos: FuncCursor) -> WasmResult<()> {

crates/environ/src/builtin.rs

@@ -45,6 +45,18 @@ macro_rules! foreach_builtin_function {
             externref_global_get(vmctx, i32) -> (reference);
             /// Returns an index for Wasm's `global.get` instruction for `externref`s.
             externref_global_set(vmctx, i32, reference) -> ();
+            /// Returns an index for wasm's `memory.atomic.notify` for locally defined memories.
+            memory_atomic_notify(vmctx, i32, i32, i32) -> (i32);
+            /// Returns an index for wasm's `memory.atomic.notify` for imported memories.
+            imported_memory_atomic_notify(vmctx, i32, i32, i32) -> (i32);
+            /// Returns an index for wasm's `memory.atomic.wait32` for locally defined memories.
+            memory_atomic_wait32(vmctx, i32, i32, i32, i64) -> (i32);
+            /// Returns an index for wasm's `memory.atomic.wait32` for imported memories.
+            imported_memory_atomic_wait32(vmctx, i32, i32, i32, i64) -> (i32);
+            /// Returns an index for wasm's `memory.atomic.wait64` for locally defined memories.
+            memory_atomic_wait64(vmctx, i32, i32, i64, i64) -> (i32);
+            /// Returns an index for wasm's `memory.atomic.wait64` for imported memories.
+            imported_memory_atomic_wait64(vmctx, i32, i32, i64, i64) -> (i32);
         }
     };
 }

crates/runtime/src/libcalls.rs

@@ -58,7 +58,7 @@
 
 use crate::externref::VMExternRef;
 use crate::table::Table;
-use crate::traphandlers::raise_lib_trap;
+use crate::traphandlers::{raise_lib_trap, Trap};
 use crate::vmcontext::{VMCallerCheckedAnyfunc, VMContext};
 use std::mem;
 use std::ptr::{self, NonNull};
@@ -496,3 +496,88 @@ pub unsafe extern "C" fn wasmtime_externref_global_set(
     let old = mem::replace((*global).as_externref_mut(), externref);
     drop(old);
 }
+
+#[derive(Debug)]
+struct Unimplemented(&'static str);
+impl std::error::Error for Unimplemented {}
+impl std::fmt::Display for Unimplemented {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::result::Result<(), std::fmt::Error> {
+        write!(f, "unimplemented: {}", self.0)
+    }
+}
+
+/// Implementation of `memory.atomic.notify` for locally defined memories.
+pub unsafe extern "C" fn wasmtime_memory_atomic_notify(
+    _vmctx: *mut VMContext,
+    _memory_index: u32,
+    _addr: u32,
+    _count: u32,
+) -> u32 {
+    raise_lib_trap(Trap::User(Box::new(Unimplemented(
+        "wasm atomics (fn wasmtime_memory_atomic_notify) unsupported",
+    ))));
+}
+
+/// Implementation of `memory.atomic.notify` for imported memories.
+pub unsafe extern "C" fn wasmtime_imported_memory_atomic_notify(
+    _vmctx: *mut VMContext,
+    _memory_index: u32,
+    _addr: u32,
+    _count: u32,
+) -> u32 {
+    raise_lib_trap(Trap::User(Box::new(Unimplemented(
+        "wasm atomics (fn wasmtime_imported_memory_atomic_notify) unsupported",
+    ))));
+}
+
+/// Implementation of `memory.atomic.wait32` for locally defined memories.
+pub unsafe extern "C" fn wasmtime_memory_atomic_wait32(
+    _vmctx: *mut VMContext,
+    _memory_index: u32,
+    _addr: u32,
+    _expected: u32,
+    _timeout: u64,
+) -> u32 {
+    raise_lib_trap(Trap::User(Box::new(Unimplemented(
+        "wasm atomics (fn wasmtime_memory_atomic_wait32) unsupported",
+    ))));
+}
+
+/// Implementation of `memory.atomic.wait32` for imported memories.
+pub unsafe extern "C" fn wasmtime_imported_memory_atomic_wait32(
+    _vmctx: *mut VMContext,
+    _memory_index: u32,
+    _addr: u32,
+    _expected: u32,
+    _timeout: u64,
+) -> u32 {
+    raise_lib_trap(Trap::User(Box::new(Unimplemented(
+        "wasm atomics (fn wasmtime_imported_memory_atomic_wait32) unsupported",
+    ))));
+}
+
+/// Implementation of `memory.atomic.wait64` for locally defined memories.
+pub unsafe extern "C" fn wasmtime_memory_atomic_wait64(
+    _vmctx: *mut VMContext,
+    _memory_index: u32,
+    _addr: u32,
+    _expected: u64,
+    _timeout: u64,
+) -> u32 {
+    raise_lib_trap(Trap::User(Box::new(Unimplemented(
+        "wasm atomics (fn wasmtime_memory_atomic_wait32) unsupported",
+    ))));
+}
+
+/// Implementation of `memory.atomic.wait32` for imported memories.
+pub unsafe extern "C" fn wasmtime_imported_memory_atomic_wait64(
+    _vmctx: *mut VMContext,
+    _memory_index: u32,
+    _addr: u32,
+    _expected: u64,
+    _timeout: u64,
+) -> u32 {
+    raise_lib_trap(Trap::User(Box::new(Unimplemented(
+        "wasm atomics (fn wasmtime_imported_memory_atomic_wait64) unsupported",
+    ))));
+}

crates/runtime/src/vmcontext.rs

@@ -600,6 +600,18 @@ impl VMBuiltinFunctionsArray {
             wasmtime_table_fill as usize;
         ptrs[BuiltinFunctionIndex::table_fill_funcref().index() as usize] =
             wasmtime_table_fill as usize;
+        ptrs[BuiltinFunctionIndex::memory_atomic_notify().index() as usize] =
+            wasmtime_memory_atomic_notify as usize;
+        ptrs[BuiltinFunctionIndex::imported_memory_atomic_notify().index() as usize] =
+            wasmtime_imported_memory_atomic_notify as usize;
+        ptrs[BuiltinFunctionIndex::memory_atomic_wait32().index() as usize] =
+            wasmtime_memory_atomic_wait32 as usize;
+        ptrs[BuiltinFunctionIndex::imported_memory_atomic_wait32().index() as usize] =
+            wasmtime_imported_memory_atomic_wait32 as usize;
+        ptrs[BuiltinFunctionIndex::memory_atomic_wait64().index() as usize] =
+            wasmtime_memory_atomic_wait64 as usize;
+        ptrs[BuiltinFunctionIndex::imported_memory_atomic_wait64().index() as usize] =
+            wasmtime_imported_memory_atomic_wait64 as usize;
 
         if cfg!(debug_assertions) {
             for i in 0..ptrs.len() {