Repository: tetratelabs/wazero
Pull request #2121 (2024-03-06 02:39:48)

wazevo(amd64): encoding for lock cmpxchg and lock xadd


Patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr.go b/internal/engine/wazevo/backend/isa/amd64/instr.go
index 74d66f5a..7b9cf1e2 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr.go
@@ -271,6 +271,32 @@ func (i *instruction) String() string {
 		return fmt.Sprintf("blendvpd %s, %s, %%xmm0", i.op1.format(false), i.op2.format(false))
 	case mfence:
 		return "mfence"
+	case lockcmpxchg:
+		var suffix string
+		switch i.u1 {
+		case 1:
+			suffix = "b"
+		case 2:
+			suffix = "w"
+		case 4:
+			suffix = "l"
+		case 8:
+			suffix = "q"
+		}
+		return fmt.Sprintf("lock cmpxchg.%s %s, %s", suffix, i.op1.format(true), i.op2.format(true))
+	case lockxadd:
+		var suffix string
+		switch i.u1 {
+		case 1:
+			suffix = "b"
+		case 2:
+			suffix = "w"
+		case 4:
+			suffix = "l"
+		case 8:
+			suffix = "q"
+		}
+		return fmt.Sprintf("lock xadd.%s %s, %s", suffix, i.op1.format(true), i.op2.format(true))
 	default:
 		panic(fmt.Sprintf("BUG: %d", int(i.kind)))
 	}
@@ -739,6 +765,9 @@ const (
 	// If the dst is a memory address, the execution is atomic.
 	xchg
 
+	// lockcmpxchg is the cmpxchg instruction https://www.felixcloutier.com/x86/cmpxchg with a lock prefix.
+	lockcmpxchg
+
 	// zeros puts zeros into the destination register. This is implemented as xor reg, reg for
 	// either integer or XMM registers. The reason why we have this instruction instead of using aluRmiR
 	// is that it requires the already-defined registers. From reg alloc's perspective, this defines
@@ -770,6 +799,9 @@ const (
 	// mfence is https://www.felixcloutier.com/x86/mfence
 	mfence
 
+	// lockxadd is xadd https://www.felixcloutier.com/x86/xadd with a lock prefix.
+	lockxadd
+
 	instrMax
 )
 
@@ -1001,6 +1033,10 @@ func (k instructionKind) String() string {
 		return "idivRemSequence"
 	case mfence:
 		return "mfence"
+	case lockcmpxchg:
+		return "lockcmpxchg"
+	case lockxadd:
+		return "lockxadd"
 	default:
 		panic("BUG")
 	}
@@ -1437,6 +1473,22 @@ func (i *instruction) asXCHG(rm regalloc.VReg, rd operand, size byte) *instructi
 	return i
 }
 
+func (i *instruction) asLockCmpXCHG(rm regalloc.VReg, rd *amode, size byte) *instruction {
+	i.kind = lockcmpxchg
+	i.op1 = newOperandReg(rm)
+	i.op2 = newOperandMem(rd)
+	i.u1 = uint64(size)
+	return i
+}
+
+func (i *instruction) asLockXAdd(rm regalloc.VReg, rd *amode, size byte) *instruction {
+	i.kind = lockxadd
+	i.op1 = newOperandReg(rm)
+	i.op2 = newOperandMem(rd)
+	i.u1 = uint64(size)
+	return i
+}
+
 type unaryRmROpcode byte
 
 const (
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go b/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
index 201f18de..80f766bf 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
@@ -1286,6 +1286,78 @@ func (i *instruction) encode(c backend.Compiler) (needsLabelResolution bool) {
 			panic("BUG: invalid operand kind")
 		}
 
+	case lockcmpxchg:
+		src, dst := regEncodings[i.op1.reg().RealReg()], i.op2
+		size := i.u1
+
+		var rex rexInfo
+		var opcode uint32
+		lp := legacyPrefixes0xF0 // Lock prefix.
+		switch size {
+		case 8:
+			opcode = 0x0FB1
+			rex = rexInfo(0).setW()
+		case 4:
+			opcode = 0x0FB1
+			rex = rexInfo(0).clearW()
+		case 2:
+			lp = legacyPrefixes0x660xF0 // Legacy prefix + Lock prefix.
+			opcode = 0x0FB1
+			rex = rexInfo(0).clearW()
+		case 1:
+			opcode = 0x0FB0
+			// Some destinations must be encoded with REX.R = 1.
+			if e := src.encoding(); e >= 4 && e <= 7 {
+				rex = rexInfo(0).always()
+			}
+		default:
+			panic(fmt.Sprintf("BUG: invalid size %d: %s", size, i.String()))
+		}
+
+		switch dst.kind {
+		case operandKindMem:
+			m := dst.addressMode()
+			encodeRegMem(c, lp, opcode, 2, src, m, rex)
+		default:
+			panic("BUG: invalid operand kind")
+		}
+
+	case lockxadd:
+		src, dst := regEncodings[i.op1.reg().RealReg()], i.op2
+		size := i.u1
+
+		var rex rexInfo
+		var opcode uint32
+		lp := legacyPrefixes0xF0 // Lock prefix.
+		switch size {
+		case 8:
+			opcode = 0x0FC1
+			rex = rexInfo(0).setW()
+		case 4:
+			opcode = 0x0FC1
+			rex = rexInfo(0).clearW()
+		case 2:
+			lp = legacyPrefixes0x660xF0 // Legacy prefix + Lock prefix.
+			opcode = 0x0FC1
+			rex = rexInfo(0).clearW()
+		case 1:
+			opcode = 0x0FC0
+			// Some destinations must be encoded with REX.R = 1.
+			if e := src.encoding(); e >= 4 && e <= 7 {
+				rex = rexInfo(0).always()
+			}
+		default:
+			panic(fmt.Sprintf("BUG: invalid size %d: %s", size, i.String()))
+		}
+
+		switch dst.kind {
+		case operandKindMem:
+			m := dst.addressMode()
+			encodeRegMem(c, lp, opcode, 2, src, m, rex)
+		default:
+			panic("BUG: invalid operand kind")
+		}
+
 	case zeros:
 		r := i.op2.reg()
 		if r.RegType() == regalloc.RegTypeInt {


Test patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr_encoding_test.go b/internal/engine/wazevo/backend/isa/amd64/instr_encoding_test.go
index 2d1b7f13..d8bb29d6 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr_encoding_test.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr_encoding_test.go
@@ -4121,6 +4121,116 @@ func TestInstruction_format_encode(t *testing.T) {
 			want:       "0faef0",
 			wantFormat: "mfence",
 		},
+		{
+			setup:      func(i *instruction) { i.asLockCmpXCHG(rsiVReg, newAmodeImmReg(123, raxVReg), 8) },
+			want:       "f0480fb1707b",
+			wantFormat: "lock cmpxchg.q %rsi, 123(%rax)",
+		},
+		{
+			setup:      func(i *instruction) { i.asLockCmpXCHG(r11VReg, newAmodeImmReg(123, raxVReg), 8) },
+			want:       "f04c0fb1587b",
+			wantFormat: "lock cmpxchg.q %r11, 123(%rax)",
+		},
+		{
+			setup:      func(i *instruction) { i.asLockCmpXCHG(rsiVReg, newAmodeImmReg(123, raxVReg), 4) },
+			want:       "f00fb1707b",
+			wantFormat: "lock cmpxchg.l %rsi, 123(%rax)",
+		},
+		{
+			setup:      func(i *instruction) { i.asLockCmpXCHG(r11VReg, newAmodeImmReg(123, raxVReg), 4) },
+			want:       "f0440fb1587b",
+			wantFormat: "lock cmpxchg.l %r11, 123(%rax)",
+		},
+		{
+			setup:      func(i *instruction) { i.asLockCmpXCHG(r11VReg, newAmodeImmReg(123, raxVReg), 2) },
+			want:       "66f0440fb1587b",
+			wantFormat: "lock cmpxchg.w %r11, 123(%rax)",
+		},
+		{
+			setup:      func(i *instruction) { i.asLockCmpXCHG(rsiVReg, newAmodeImmReg(123, raxVReg), 2) },
+			want:       "66f00fb1707b",
+			wantFormat: "lock cmpxchg.w %rsi, 123(%rax)",
+		},
+		{
+			setup:      func(i *instruction) { i.asLockCmpXCHG(r11VReg, newAmodeImmReg(123, raxVReg), 1) },
+			want:       "f0440fb0587b",
+			wantFormat: "lock cmpxchg.b %r11, 123(%rax)",
+		},
+		{
+			setup:      func(i *instruction) { i.asLockCmpXCHG(r11VReg, newAmodeImmReg(123, rdiVReg), 1) },
+			want:       "f0440fb05f7b",
+			wantFormat: "lock cmpxchg.b %r11, 123(%rdi)",
+		},
+		{
+			setup:      func(i *instruction) { i.asLockCmpXCHG(rsiVReg, newAmodeImmReg(123, raxVReg), 1) },
+			want:       "f0400fb0707b",
+			wantFormat: "lock cmpxchg.b %rsi, 123(%rax)",
+		},
+		{
+			setup:      func(i *instruction) { i.asLockCmpXCHG(rdiVReg, newAmodeImmReg(123, raxVReg), 1) },
+			want:       "f0400fb0787b",
+			wantFormat: "lock cmpxchg.b %rdi, 123(%rax)",
+		},
+		{
+			setup:      func(i *instruction) { i.asLockCmpXCHG(rdiVReg, newAmodeImmReg(123, rsiVReg), 1) },
+			want:       "f0400fb07e7b",
+			wantFormat: "lock cmpxchg.b %rdi, 123(%rsi)",
+		},
+		{
+			setup:      func(i *instruction) { i.asLockXAdd(rsiVReg, newAmodeImmReg(123, raxVReg), 8) },
+			want:       "f0480fc1707b",
+			wantFormat: "lock xadd.q %rsi, 123(%rax)",
+		},
+		{
+			setup:      func(i *instruction) { i.asLockXAdd(r11VReg, newAmodeImmReg(123, raxVReg), 8) },
+			want:       "f04c0fc1587b",
+			wantFormat: "lock xadd.q %r11, 123(%rax)",
+		},
+		{
+			setup:      func(i *instruction) { i.asLockXAdd(rsiVReg, newAmodeImmReg(123, raxVReg), 4) },
+			want:       "f00fc1707b",
+			wantFormat: "lock xadd.l %rsi, 123(%rax)",
+		},
+		{
+			setup:      func(i *instruction) { i.asLockXAdd(r11VReg, newAmodeImmReg(123, raxVReg), 4) },
+			want:       "f0440fc1587b",
+			wantFormat: "lock xadd.l %r11, 123(%rax)",
+		},
+		{
+			setup:      func(i *instruction) { i.asLockXAdd(r11VReg, newAmodeImmReg(123, raxVReg), 2) },
+			want:       "66f0440fc1587b",
+			wantFormat: "lock xadd.w %r11, 123(%rax)",
+		},
+		{
+			setup:      func(i *instruction) { i.asLockXAdd(rsiVReg, newAmodeImmReg(123, raxVReg), 2) },
+			want:       "66f00fc1707b",
+			wantFormat: "lock xadd.w %rsi, 123(%rax)",
+		},
+		{
+			setup:      func(i *instruction) { i.asLockXAdd(r11VReg, newAmodeImmReg(123, raxVReg), 1) },
+			want:       "f0440fc0587b",
+			wantFormat: "lock xadd.b %r11, 123(%rax)",
+		},
+		{
+			setup:      func(i *instruction) { i.asLockXAdd(r11VReg, newAmodeImmReg(123, rdiVReg), 1) },
+			want:       "f0440fc05f7b",
+			wantFormat: "lock xadd.b %r11, 123(%rdi)",
+		},
+		{
+			setup:      func(i *instruction) { i.asLockXAdd(rsiVReg, newAmodeImmReg(123, raxVReg), 1) },
+			want:       "f0400fc0707b",
+			wantFormat: "lock xadd.b %rsi, 123(%rax)",
+		},
+		{
+			setup:      func(i *instruction) { i.asLockXAdd(rdiVReg, newAmodeImmReg(123, raxVReg), 1) },
+			want:       "f0400fc0787b",
+			wantFormat: "lock xadd.b %rdi, 123(%rax)",
+		},
+		{
+			setup:      func(i *instruction) { i.asLockXAdd(rdiVReg, newAmodeImmReg(123, rsiVReg), 1) },
+			want:       "f0400fc07e7b",
+			wantFormat: "lock xadd.b %rdi, 123(%rsi)",
+		},
 	} {
 		tc := tc
 		t.Run(tc.wantFormat, func(t *testing.T) {