Repository: tetratelabs/wazero
Pull request #2123 (2024-03-06 04:39:52)

wazevo(amd64): atomic rmw add/sub/xchg


Patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr.go b/internal/engine/wazevo/backend/isa/amd64/instr.go
index 63a7d6cd..f0d272eb 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr.go
@@ -2351,6 +2351,8 @@ var defKinds = [instrMax]defKind{
 	mfence:                 defKindNone,
 	xchg:                   defKindNone,
 	lockcmpxchg:            defKindNone,
+	lockxadd:               defKindNone,
+	neg:                    defKindNone,
 }
 
 // String implements fmt.Stringer.
@@ -2431,6 +2433,8 @@ var useKinds = [instrMax]useKind{
 	mfence:                 useKindNone,
 	xchg:                   useKindOp1RegOp2,
 	lockcmpxchg:            useKindRaxOp1RegOp2,
+	lockxadd:               useKindOp1RegOp2,
+	neg:                    useKindOp1,
 }
 
 func (u useKind) String() string {
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine.go b/internal/engine/wazevo/backend/isa/amd64/machine.go
index 44a1f0ff..1f5728d4 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine.go
@@ -976,11 +976,60 @@ func (m *machine) LowerInstr(instr *ssa.Instruction) {
 		size := instr.AtomicTargetSize()
 		m.lowerAtomicCas(addr, exp, repl, size, instr.Return())
 
+	case ssa.OpcodeAtomicRmw:
+		addr, val := instr.Arg2()
+		atomicOp, size := instr.AtomicRmwData()
+		m.lowerAtomicRmw(atomicOp, addr, val, size, instr.Return())
+
 	default:
 		panic("TODO: lowering " + op.String())
 	}
 }
 
+func (m *machine) lowerAtomicRmw(op ssa.AtomicRmwOp, addr, val ssa.Value, size uint64, ret ssa.Value) {
+	mem := m.lowerToAddressMode(addr, 0)
+	_val := m.getOperand_Reg(m.c.ValueDefinition(val))
+	valCopied := m.copyToTmp(_val.reg())
+
+	switch op {
+	case ssa.AtomicRmwOpAdd, ssa.AtomicRmwOpSub:
+		if op == ssa.AtomicRmwOpSub {
+			// Negate the value.
+			m.insert(m.allocateInstr().asNeg(newOperandReg(valCopied), true))
+		}
+		m.insert(m.allocateInstr().asLockXAdd(valCopied, mem, byte(size)))
+		if size < 4 {
+			// Clear the unnecessary bits.
+			m.insert(m.allocateInstr().asAluRmiR(aluRmiROpcodeAnd, newOperandImm32(uint32((1<<(8*size))-1)),
+				valCopied, true))
+		}
+
+		m.copyTo(valCopied, m.c.VRegOf(ret))
+
+	case ssa.AtomicRmwOpAnd:
+		panic("TODO")
+
+	case ssa.AtomicRmwOpOr:
+		panic("TODO")
+
+	case ssa.AtomicRmwOpXor:
+		panic("TODO")
+
+	case ssa.AtomicRmwOpXchg:
+		m.insert(m.allocateInstr().asXCHG(valCopied, newOperandMem(mem), byte(size)))
+		if size < 4 {
+			// Clear the unnecessary bits.
+			m.insert(m.allocateInstr().asAluRmiR(aluRmiROpcodeAnd, newOperandImm32(uint32((1<<(8*size))-1)),
+				valCopied, true))
+		}
+
+		m.copyTo(valCopied, m.c.VRegOf(ret))
+
+	default:
+		panic("BUG")
+	}
+}
+
 func (m *machine) lowerAtomicCas(addr, exp, repl ssa.Value, size uint64, ret ssa.Value) {
 	mem := m.lowerToAddressMode(addr, 0)
 	expOp := m.getOperand_Reg(m.c.ValueDefinition(exp))
diff --git a/internal/engine/wazevo/ssa/instructions.go b/internal/engine/wazevo/ssa/instructions.go
index 5d2ddad0..b49fdd16 100644
--- a/internal/engine/wazevo/ssa/instructions.go
+++ b/internal/engine/wazevo/ssa/instructions.go
@@ -547,10 +547,6 @@ const (
 	// OpcodeBitcast is a bitcast operation: `v = bitcast x`.
 	OpcodeBitcast
 
-	// OpcodeBmask ...
-	// `v = bmask x`.
-	OpcodeBmask
-
 	// OpcodeIreduce narrow the given integer: `v = Ireduce x`.
 	OpcodeIreduce
 
@@ -599,18 +595,16 @@ const (
 	// to two lower single-precision lanes of the result `v = Fvdemote.lane x`.
 	OpcodeFvdemote
 
-	// OpcodeFcvtToUint ...
-	// `v = fcvt_to_uint x`.
+	// OpcodeFcvtToUint converts a floating point value to an unsigned integer: `v = FcvtToUint x`.
 	OpcodeFcvtToUint
 
 	// OpcodeFcvtToSint converts a floating point value to a signed integer: `v = FcvtToSint x`.
 	OpcodeFcvtToSint
 
-	// OpcodeFcvtToUintSat converts a floating point value to an unsigned integer: `v = FcvtToUintSat x`.
+	// OpcodeFcvtToUintSat converts a floating point value to an unsigned integer: `v = FcvtToUintSat x` which saturates on overflow.
 	OpcodeFcvtToUintSat
 
-	// OpcodeFcvtToSintSat ...
-	// `v = fcvt_to_sint_sat x`.
+	// OpcodeFcvtToSintSat converts a floating point value to a signed integer: `v = FcvtToSintSat x` which saturates on overflow.
 	OpcodeFcvtToSintSat
 
 	// OpcodeFcvtFromUint converts an unsigned integer to a floating point value: `v = FcvtFromUint x`.
@@ -619,38 +613,26 @@ const (
 	// OpcodeFcvtFromSint converts a signed integer to a floating point value: `v = FcvtFromSint x`.
 	OpcodeFcvtFromSint
 
-	// OpcodeIsplit ...
-	// `lo, hi = isplit x`.
-	OpcodeIsplit
-
-	// OpcodeIconcat ...
-	// `v = iconcat lo, hi`.
-	OpcodeIconcat
-
-	// OpcodeAtomicRmw ...
-	// `v = atomic_rmw MemFlags, AtomicRmwOp, p, x`.
+	// OpcodeAtomicRmw is atomic read-modify-write operation: `v = atomic_rmw op, p, offset, value`.
 	OpcodeAtomicRmw
 
-	// OpcodeAtomicCas ...
-	// `v = atomic_cas MemFlags, p, e, x`.
+	// OpcodeAtomicCas is atomic compare-and-swap operation.
 	OpcodeAtomicCas
 
-	// OpcodeAtomicLoad ...
-	// `v = atomic_load MemFlags, p`.
+	// OpcodeAtomicLoad is atomic load operation.
 	OpcodeAtomicLoad
 
-	// OpcodeAtomicStore ...
-	// `atomic_store MemFlags, x, p`.
+	// OpcodeAtomicStore is atomic store operation.
 	OpcodeAtomicStore
 
-	// OpcodeFence ...
-	// `fence`.
+	// OpcodeFence is a memory fence operation.
 	OpcodeFence
 
 	// opcodeEnd marks the end of the opcode list.
 	opcodeEnd
 )
 
+// AtomicRmwOp represents the atomic read-modify-write operation.
 type AtomicRmwOp byte
 
 const (
@@ -668,6 +650,7 @@ const (
 	AtomicRmwOpXchg
 )
 
+// String implements the fmt.Stringer.
 func (op AtomicRmwOp) String() string {
 	switch op {
 	case AtomicRmwOpAdd:
@@ -2829,8 +2812,6 @@ func (o Opcode) String() (ret string) {
 		return "Nearest"
 	case OpcodeBitcast:
 		return "Bitcast"
-	case OpcodeBmask:
-		return "Bmask"
 	case OpcodeIreduce:
 		return "Ireduce"
 	case OpcodeSnarrow:
@@ -2871,10 +2852,6 @@ func (o Opcode) String() (ret string) {
 		return "FcvtFromUint"
 	case OpcodeFcvtFromSint:
 		return "FcvtFromSint"
-	case OpcodeIsplit:
-		return "Isplit"
-	case OpcodeIconcat:
-		return "Iconcat"
 	case OpcodeAtomicRmw:
 		return "AtomicRmw"
 	case OpcodeAtomicCas:


Test patch:
diff --git a/internal/engine/wazevo/e2e_test.go b/internal/engine/wazevo/e2e_test.go
index 56c7ffcc..21dadb86 100644
--- a/internal/engine/wazevo/e2e_test.go
+++ b/internal/engine/wazevo/e2e_test.go
@@ -644,10 +644,9 @@ func TestE2E(t *testing.T) {
 			},
 		},
 		{
-			name:      "atomic_rmw_add",
-			m:         testcases.AtomicRmwAdd.Module,
-			skipAMD64: true,
-			features:  api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
+			name:     "atomic_rmw_add",
+			m:        testcases.AtomicRmwAdd.Module,
+			features: api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
 			calls: []callCase{
 				{params: []uint64{1, 2, 3, 4, 5, 6, 7}, expResults: []uint64{0, 0, 0, 0, 0, 0, 0}},
 				{params: []uint64{1, 2, 3, 4, 5, 6, 7}, expResults: []uint64{1, 2, 3, 4, 5, 6, 7}},
@@ -655,10 +654,9 @@ func TestE2E(t *testing.T) {
 			},
 		},
 		{
-			name:      "atomic_rmw_sub",
-			m:         testcases.AtomicRmwSub.Module,
-			features:  api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
-			skipAMD64: true,
+			name:     "atomic_rmw_sub",
+			m:        testcases.AtomicRmwSub.Module,
+			features: api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
 			calls: []callCase{
 				{params: []uint64{1, 2, 3, 4, 5, 6, 7}, expResults: []uint64{0, 0, 0, 0, 0, 0, 0}},
 				{
@@ -761,10 +759,9 @@ func TestE2E(t *testing.T) {
 			},
 		},
 		{
-			name:      "atomic_rmw_xchg",
-			m:         testcases.AtomicRmwXchg.Module,
-			features:  api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
-			skipAMD64: true,
+			name:     "atomic_rmw_xchg",
+			m:        testcases.AtomicRmwXchg.Module,
+			features: api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
 			calls: []callCase{
 				{
 					params:     []uint64{1, 2, 3, 4, 5, 6, 7},