Repository: tetratelabs/wazero
Pull request #1899 (2024-01-05 04:44:28)

threads: add wazeroir signatures for atomic ops


Patch:
diff --git a/internal/wazeroir/signature.go b/internal/wazeroir/signature.go
index 6d0f8444..30a5868d 100644
--- a/internal/wazeroir/signature.go
+++ b/internal/wazeroir/signature.go
@@ -233,6 +233,26 @@ var (
 		in:  []UnsignedType{UnsignedTypeF64},
 		out: []UnsignedType{UnsignedTypeV128},
 	}
+	signature_I32I64_I64 = &signature{
+		in:  []UnsignedType{UnsignedTypeI32, UnsignedTypeI64},
+		out: []UnsignedType{UnsignedTypeI64},
+	}
+	signature_I32I32I64_I32 = &signature{
+		in:  []UnsignedType{UnsignedTypeI32, UnsignedTypeI32, UnsignedTypeI64},
+		out: []UnsignedType{UnsignedTypeI32},
+	}
+	signature_I32I64I64_I32 = &signature{
+		in:  []UnsignedType{UnsignedTypeI32, UnsignedTypeI64, UnsignedTypeI64},
+		out: []UnsignedType{UnsignedTypeI32},
+	}
+	signature_I32I32I32_I32 = &signature{
+		in:  []UnsignedType{UnsignedTypeI32, UnsignedTypeI32, UnsignedTypeI32},
+		out: []UnsignedType{UnsignedTypeI32},
+	}
+	signature_I32I64I64_I64 = &signature{
+		in:  []UnsignedType{UnsignedTypeI32, UnsignedTypeI64, UnsignedTypeI64},
+		out: []UnsignedType{UnsignedTypeI64},
+	}
 )
 
 // wasmOpcodeSignature returns the signature of given Wasm opcode.
@@ -586,6 +606,40 @@ func (c *Compiler) wasmOpcodeSignature(op wasm.Opcode, index uint32) (*signature
 		default:
 			return nil, fmt.Errorf("unsupported vector instruction in wazeroir: %s", wasm.VectorInstructionName(vecOp))
 		}
+	case wasm.OpcodeAtomicPrefix:
+		switch atomicOp := c.body[c.pc+1]; atomicOp {
+		case wasm.OpcodeAtomicMemoryNotify:
+			return signature_I32I32_I32, nil
+		case wasm.OpcodeAtomicMemoryWait32:
+			return signature_I32I32I64_I32, nil
+		case wasm.OpcodeAtomicMemoryWait64:
+			return signature_I32I64I64_I32, nil
+		case wasm.OpcodeAtomicFence:
+			return signature_None_None, nil
+		case wasm.OpcodeAtomicI32Load, wasm.OpcodeAtomicI32Load8U, wasm.OpcodeAtomicI32Load16U:
+			return signature_I32_I32, nil
+		case wasm.OpcodeAtomicI64Load, wasm.OpcodeAtomicI64Load8U, wasm.OpcodeAtomicI64Load16U, wasm.OpcodeAtomicI64Load32U:
+			return signature_I32_I64, nil
+		case wasm.OpcodeAtomicI32Store, wasm.OpcodeAtomicI32Store8, wasm.OpcodeAtomicI32Store16:
+			return signature_I32I32_None, nil
+		case wasm.OpcodeAtomicI64Store, wasm.OpcodeAtomicI64Store8, wasm.OpcodeAtomicI64Store16, wasm.OpcodeAtomicI64Store32:
+			return signature_I32I64_None, nil
+		case wasm.OpcodeAtomicI32RmwAdd, wasm.OpcodeAtomicI32RmwSub, wasm.OpcodeAtomicI32RmwAnd, wasm.OpcodeAtomicI32RmwOr, wasm.OpcodeAtomicI32RmwXor, wasm.OpcodeAtomicI32RmwXchg,
+			wasm.OpcodeAtomicI32Rmw8AddU, wasm.OpcodeAtomicI32Rmw8SubU, wasm.OpcodeAtomicI32Rmw8AndU, wasm.OpcodeAtomicI32Rmw8OrU, wasm.OpcodeAtomicI32Rmw8XorU, wasm.OpcodeAtomicI32Rmw8XchgU,
+			wasm.OpcodeAtomicI32Rmw16AddU, wasm.OpcodeAtomicI32Rmw16SubU, wasm.OpcodeAtomicI32Rmw16AndU, wasm.OpcodeAtomicI32Rmw16OrU, wasm.OpcodeAtomicI32Rmw16XorU, wasm.OpcodeAtomicI32Rmw16XchgU:
+			return signature_I32I32_I32, nil
+		case wasm.OpcodeAtomicI64RmwAdd, wasm.OpcodeAtomicI64RmwSub, wasm.OpcodeAtomicI64RmwAnd, wasm.OpcodeAtomicI64RmwOr, wasm.OpcodeAtomicI64RmwXor, wasm.OpcodeAtomicI64RmwXchg,
+			wasm.OpcodeAtomicI64Rmw8AddU, wasm.OpcodeAtomicI64Rmw8SubU, wasm.OpcodeAtomicI64Rmw8AndU, wasm.OpcodeAtomicI64Rmw8OrU, wasm.OpcodeAtomicI64Rmw8XorU, wasm.OpcodeAtomicI64Rmw8XchgU,
+			wasm.OpcodeAtomicI64Rmw16AddU, wasm.OpcodeAtomicI64Rmw16SubU, wasm.OpcodeAtomicI64Rmw16AndU, wasm.OpcodeAtomicI64Rmw16OrU, wasm.OpcodeAtomicI64Rmw16XorU, wasm.OpcodeAtomicI64Rmw16XchgU,
+			wasm.OpcodeAtomicI64Rmw32AddU, wasm.OpcodeAtomicI64Rmw32SubU, wasm.OpcodeAtomicI64Rmw32AndU, wasm.OpcodeAtomicI64Rmw32OrU, wasm.OpcodeAtomicI64Rmw32XorU, wasm.OpcodeAtomicI64Rmw32XchgU:
+			return signature_I32I64_I64, nil
+		case wasm.OpcodeAtomicI32RmwCmpxchg, wasm.OpcodeAtomicI32Rmw8CmpxchgU, wasm.OpcodeAtomicI32Rmw16CmpxchgU:
+			return signature_I32I32I32_I32, nil
+		case wasm.OpcodeAtomicI64RmwCmpxchg, wasm.OpcodeAtomicI64Rmw8CmpxchgU, wasm.OpcodeAtomicI64Rmw16CmpxchgU, wasm.OpcodeAtomicI64Rmw32CmpxchgU:
+			return signature_I32I64I64_I64, nil
+		default:
+			return nil, fmt.Errorf("unsupported atomic instruction in wazeroir: %s", wasm.AtomicInstructionName(atomicOp))
+		}
 	default:
 		return nil, fmt.Errorf("unsupported instruction in wazeroir: 0x%x", op)
 	}


Test patch:
diff --git a/internal/wazeroir/signature_test.go b/internal/wazeroir/signature_test.go
index 3969de7c..51867836 100644
--- a/internal/wazeroir/signature_test.go
+++ b/internal/wazeroir/signature_test.go
@@ -88,6 +88,341 @@ func TestCompiler_wasmOpcodeSignature(t *testing.T) {
 			body: []byte{wasm.OpcodeMiscPrefix, wasm.OpcodeMiscTableCopy},
 			exp:  signature_I32I32I32_None,
 		},
+		{
+			name: "i32.atomic.load8_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32Load8U},
+			exp:  signature_I32_I32,
+		},
+		{
+			name: "i32.atomic.load16_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32Load16U},
+			exp:  signature_I32_I32,
+		},
+		{
+			name: "i32.atomic.load",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32Load},
+			exp:  signature_I32_I32,
+		},
+		{
+			name: "i64.atomic.load8_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Load8U},
+			exp:  signature_I32_I64,
+		},
+		{
+			name: "i64.atomic.load16_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Load16U},
+			exp:  signature_I32_I64,
+		},
+		{
+			name: "i64.atomic.load32_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Load32U},
+			exp:  signature_I32_I64,
+		},
+		{
+			name: "i64.atomic.load",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Load},
+			exp:  signature_I32_I64,
+		},
+		{
+			name: "i32.atomic.store8",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32Store8},
+			exp:  signature_I32I32_None,
+		},
+		{
+			name: "i32.atomic.store16_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32Store16},
+			exp:  signature_I32I32_None,
+		},
+		{
+			name: "i32.atomic.store",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32Store},
+			exp:  signature_I32I32_None,
+		},
+		{
+			name: "i64.atomic.store8",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Store8},
+			exp:  signature_I32I64_None,
+		},
+		{
+			name: "i64.atomic.store16",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Store16},
+			exp:  signature_I32I64_None,
+		},
+		{
+			name: "i64.atomic.store32",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Store32},
+			exp:  signature_I32I64_None,
+		},
+		{
+			name: "i64.atomic.store",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Store},
+			exp:  signature_I32I64_None,
+		},
+		{
+			name: "i32.atomic.rmw8.add_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32Rmw8AddU},
+			exp:  signature_I32I32_I32,
+		},
+		{
+			name: "i32.atomic.rmw16.add_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32Rmw16AddU},
+			exp:  signature_I32I32_I32,
+		},
+		{
+			name: "i32.atomic.rmw.add",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32RmwAdd},
+			exp:  signature_I32I32_I32,
+		},
+		{
+			name: "i64.atomic.rmw8.add_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Rmw8AddU},
+			exp:  signature_I32I64_I64,
+		},
+		{
+			name: "i64.atomic.rmw16.add_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Rmw16AddU},
+			exp:  signature_I32I64_I64,
+		},
+		{
+			name: "i64.atomic.rmw32.add_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Rmw32AddU},
+			exp:  signature_I32I64_I64,
+		},
+		{
+			name: "i64.atomic.rmw.add",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64RmwAdd},
+			exp:  signature_I32I64_I64,
+		},
+		{
+			name: "i32.atomic.rmw8.sub_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32Rmw8SubU},
+			exp:  signature_I32I32_I32,
+		},
+		{
+			name: "i32.atomic.rmw16.sub_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32Rmw16SubU},
+			exp:  signature_I32I32_I32,
+		},
+		{
+			name: "i32.atomic.rmw.sub",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32RmwSub},
+			exp:  signature_I32I32_I32,
+		},
+		{
+			name: "i64.atomic.rmw8.sub_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Rmw8SubU},
+			exp:  signature_I32I64_I64,
+		},
+		{
+			name: "i64.atomic.rmw16.sub_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Rmw16SubU},
+			exp:  signature_I32I64_I64,
+		},
+		{
+			name: "i64.atomic.rmw32.sub_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Rmw32SubU},
+			exp:  signature_I32I64_I64,
+		},
+		{
+			name: "i64.atomic.rmw.sub",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64RmwSub},
+			exp:  signature_I32I64_I64,
+		},
+		{
+			name: "i32.atomic.rmw8.and_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32Rmw8AndU},
+			exp:  signature_I32I32_I32,
+		},
+		{
+			name: "i32.atomic.rmw16.and_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32Rmw16AndU},
+			exp:  signature_I32I32_I32,
+		},
+		{
+			name: "i32.atomic.rmw.and",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32RmwAnd},
+			exp:  signature_I32I32_I32,
+		},
+		{
+			name: "i64.atomic.rmw8.and_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Rmw8AndU},
+			exp:  signature_I32I64_I64,
+		},
+		{
+			name: "i64.atomic.rmw16.and_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Rmw16AndU},
+			exp:  signature_I32I64_I64,
+		},
+		{
+			name: "i64.atomic.rmw32.and_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Rmw32AndU},
+			exp:  signature_I32I64_I64,
+		},
+		{
+			name: "i64.atomic.rmw.and",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64RmwAnd},
+			exp:  signature_I32I64_I64,
+		},
+		{
+			name: "i32.atomic.rmw8.or_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32Rmw8OrU},
+			exp:  signature_I32I32_I32,
+		},
+		{
+			name: "i32.atomic.rmw16.or_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32Rmw16OrU},
+			exp:  signature_I32I32_I32,
+		},
+		{
+			name: "i32.atomic.rmw.or",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32RmwOr},
+			exp:  signature_I32I32_I32,
+		},
+		{
+			name: "i64.atomic.rmw8.or_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Rmw8OrU},
+			exp:  signature_I32I64_I64,
+		},
+		{
+			name: "i64.atomic.rmw16.or_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Rmw16OrU},
+			exp:  signature_I32I64_I64,
+		},
+		{
+			name: "i64.atomic.rmw32.or_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Rmw32OrU},
+			exp:  signature_I32I64_I64,
+		},
+		{
+			name: "i64.atomic.rmw.or",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64RmwOr},
+			exp:  signature_I32I64_I64,
+		},
+		{
+			name: "i32.atomic.rmw8.xor_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32Rmw8XorU},
+			exp:  signature_I32I32_I32,
+		},
+		{
+			name: "i32.atomic.rmw16.xor_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32Rmw16XorU},
+			exp:  signature_I32I32_I32,
+		},
+		{
+			name: "i32.atomic.rmw.xor",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32RmwXor},
+			exp:  signature_I32I32_I32,
+		},
+		{
+			name: "i64.atomic.rmw8.xor_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Rmw8XorU},
+			exp:  signature_I32I64_I64,
+		},
+		{
+			name: "i64.atomic.rmw16.xor_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Rmw16XorU},
+			exp:  signature_I32I64_I64,
+		},
+		{
+			name: "i64.atomic.rmw32.xor_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Rmw32XorU},
+			exp:  signature_I32I64_I64,
+		},
+		{
+			name: "i64.atomic.rmw.xor",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64RmwXor},
+			exp:  signature_I32I64_I64,
+		},
+		{
+			name: "i32.atomic.rmw8.xchg_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32Rmw8XchgU},
+			exp:  signature_I32I32_I32,
+		},
+		{
+			name: "i32.atomic.rmw16.xchg_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32Rmw16XchgU},
+			exp:  signature_I32I32_I32,
+		},
+		{
+			name: "i32.atomic.rmw.xchg",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32RmwXchg},
+			exp:  signature_I32I32_I32,
+		},
+		{
+			name: "i64.atomic.rmw8.xchg_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Rmw8XchgU},
+			exp:  signature_I32I64_I64,
+		},
+		{
+			name: "i64.atomic.rmw16.xchg_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Rmw16XchgU},
+			exp:  signature_I32I64_I64,
+		},
+		{
+			name: "i64.atomic.rmw32.xchg_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Rmw32XchgU},
+			exp:  signature_I32I64_I64,
+		},
+		{
+			name: "i64.atomic.rmw.xchg",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64RmwXchg},
+			exp:  signature_I32I64_I64,
+		},
+		{
+			name: "i32.atomic.rmw8.cmpxchg_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32Rmw8CmpxchgU},
+			exp:  signature_I32I32I32_I32,
+		},
+		{
+			name: "i32.atomic.rmw16.cmpxchg_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32Rmw16CmpxchgU},
+			exp:  signature_I32I32I32_I32,
+		},
+		{
+			name: "i32.atomic.rmw.cmpxchg",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32RmwCmpxchg},
+			exp:  signature_I32I32I32_I32,
+		},
+		{
+			name: "i64.atomic.rmw8.cmpxchg_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Rmw8CmpxchgU},
+			exp:  signature_I32I64I64_I64,
+		},
+		{
+			name: "i64.atomic.rmw16.cmpxchg_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Rmw16CmpxchgU},
+			exp:  signature_I32I64I64_I64,
+		},
+		{
+			name: "i64.atomic.rmw32.cmpxchg_u",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Rmw32CmpxchgU},
+			exp:  signature_I32I64I64_I64,
+		},
+		{
+			name: "i64.atomic.rmw.cmpxchg",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64RmwCmpxchg},
+			exp:  signature_I32I64I64_I64,
+		},
+		{
+			name: "memory.atomic.wait32",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicMemoryWait32},
+			exp:  signature_I32I32I64_I32,
+		},
+		{
+			name: "memory.atomic.wait64",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicMemoryWait64},
+			exp:  signature_I32I64I64_I32,
+		},
+		{
+			name: "memory.atomic.notify",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicMemoryNotify},
+			exp:  signature_I32I32_I32,
+		},
+		{
+			name: "memory.atomic.fence",
+			body: []byte{wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicFence},
+			exp:  signature_None_None,
+		},
 	}
 
 	for _, tt := range tests {