wazevo(amd64): support for atomic CAS (#2122)

Signed-off-by: Takeshi Yoneda <[email protected]>
tetratelabs · Mar 6, 2024 · 55d389f · 55d389f
1 parent e594899
commit 55d389f
Show file tree

Hide file tree

Showing 3 changed files with 105 additions and 27 deletions.
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr.go b/internal/engine/wazevo/backend/isa/amd64/instr.go
@@ -411,6 +411,22 @@ func (i *instruction) Uses(regs *[]regalloc.VReg) []regalloc.VReg {
 			panic(fmt.Sprintf("BUG: invalid operand: %s", i))
 		}
 		*regs = append(*regs, opReg.reg())
+
+	case useKindRaxOp1RegOp2:
+		opReg, opAny := &i.op1, &i.op2
+		*regs = append(*regs, raxVReg, opReg.reg())
+		switch opAny.kind {
+		case operandKindReg:
+			*regs = append(*regs, opAny.reg())
+		case operandKindMem:
+			opAny.addressMode().uses(regs)
+		default:
+			panic(fmt.Sprintf("BUG: invalid operand: %s", i))
+		}
+		if opReg.kind != operandKindReg {
+			panic(fmt.Sprintf("BUG: invalid operand: %s", i))
+		}
+
 	default:
 		panic(fmt.Sprintf("BUG: invalid useKind %s for %s", uk, i))
 	}
@@ -532,31 +548,69 @@ func (i *instruction) AssignUse(index int, v regalloc.VReg) {
 		}
 	case useKindBlendvpd:
 		op, opMustBeReg := &i.op1, &i.op2
-		switch op.kind {
-		case operandKindReg:
-			switch index {
-			case 0:
-				if v.RealReg() != xmm0 {
+		if index == 0 {
+			if v.RealReg() != xmm0 {
+				panic("BUG")
+			}
+		} else {
+			switch op.kind {
+			case operandKindReg:
+				switch index {
+				case 1:
+					op.setReg(v)
+				case 2:
+					opMustBeReg.setReg(v)
+				default:
+					panic("BUG")
+				}
+			case operandKindMem:
+				nregs := op.addressMode().nregs()
+				index--
+				if index < nregs {
+					op.addressMode().assignUses(index, v)
+				} else if index == nregs {
+					opMustBeReg.setReg(v)
+				} else {
 					panic("BUG")
 				}
-			case 1:
-				op.setReg(v)
-			case 2:
-				opMustBeReg.setReg(v)
 			default:
-				panic("BUG")
+				panic(fmt.Sprintf("BUG: invalid operand pair: %s", i))
 			}
-		case operandKindMem:
-			nregs := op.addressMode().nregs()
-			if index < nregs {
-				op.addressMode().assignUses(index, v)
-			} else if index == nregs {
-				opMustBeReg.setReg(v)
-			} else {
+		}
+
+	case useKindRaxOp1RegOp2:
+		switch index {
+		case 0:
+			if v.RealReg() != rax {
 				panic("BUG")
 			}
+		case 1:
+			i.op1.setReg(v)
 		default:
-			panic(fmt.Sprintf("BUG: invalid operand pair: %s", i))
+			op := &i.op2
+			switch op.kind {
+			case operandKindReg:
+				switch index {
+				case 1:
+					op.setReg(v)
+				case 2:
+					op.setReg(v)
+				default:
+					panic("BUG")
+				}
+			case operandKindMem:
+				nregs := op.addressMode().nregs()
+				index -= 2
+				if index < nregs {
+					op.addressMode().assignUses(index, v)
+				} else if index == nregs {
+					op.setReg(v)
+				} else {
+					panic("BUG")
+				}
+			default:
+				panic(fmt.Sprintf("BUG: invalid operand pair: %s", i))
+			}
 		}
 	default:
 		panic(fmt.Sprintf("BUG: invalid useKind %s for %s", uk, i))
@@ -2296,6 +2350,7 @@ var defKinds = [instrMax]defKind{
 	blendvpd:               defKindNone,
 	mfence:                 defKindNone,
 	xchg:                   defKindNone,
+	lockcmpxchg:            defKindNone,
 }
 
 // String implements fmt.Stringer.
@@ -2323,6 +2378,8 @@ const (
 	useKindOp1Op2Reg
 	// useKindOp1RegOp2 is Op1 must be a register, Op2 can be any operand.
 	useKindOp1RegOp2
+	// useKindRaxOp1RegOp2 is Op1 must be a register, Op2 can be any operand, and RAX is used.
+	useKindRaxOp1RegOp2
 	useKindDivRem
 	useKindBlendvpd
 	useKindCall
@@ -2373,6 +2430,7 @@ var useKinds = [instrMax]useKind{
 	blendvpd:               useKindBlendvpd,
 	mfence:                 useKindNone,
 	xchg:                   useKindOp1RegOp2,
+	lockcmpxchg:            useKindRaxOp1RegOp2,
 }
 
 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
@@ -971,11 +971,33 @@ func (m *machine) LowerInstr(instr *ssa.Instruction) {
 		store := m.allocateInstr().asXCHG(copied, mem, byte(size))
 		m.insert(store)
 
+	case ssa.OpcodeAtomicCas:
+		addr, exp, repl := instr.Arg3()
+		size := instr.AtomicTargetSize()
+		m.lowerAtomicCas(addr, exp, repl, size, instr.Return())
+
 	default:
 		panic("TODO: lowering " + op.String())
 	}
 }
 
+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))
+	replOp := m.getOperand_Reg(m.c.ValueDefinition(repl))
+
+	accumulator := raxVReg
+	m.copyTo(expOp.reg(), accumulator)
+	m.insert(m.allocateInstr().asLockCmpXCHG(replOp.reg(), mem, byte(size)))
+
+	if size < 4 {
+		// Clear the unnecessary bits.
+		m.insert(m.allocateInstr().asAluRmiR(aluRmiROpcodeAnd, newOperandImm32(uint32((1<<(8*size))-1)), accumulator, true))
+	}
+
+	m.copyTo(accumulator, m.c.VRegOf(ret))
+}
+
 func (m *machine) lowerFcmp(instr *ssa.Instruction) {
 	f1, f2, and := m.lowerFcmpToFlags(instr)
 	rd := m.c.VRegOf(instr.Return())

diff --git a/internal/engine/wazevo/e2e_test.go b/internal/engine/wazevo/e2e_test.go
@@ -646,8 +646,8 @@ func TestE2E(t *testing.T) {
 		{
 			name:      "atomic_rmw_add",
 			m:         testcases.AtomicRmwAdd.Module,
-			features:  api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
 			skipAMD64: true,
+			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}},
@@ -781,10 +781,9 @@ func TestE2E(t *testing.T) {
 			},
 		},
 		{
-			name:      "atomic_cas",
-			m:         testcases.AtomicCas.Module,
-			features:  api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
-			skipAMD64: true,
+			name:     "atomic_cas",
+			m:        testcases.AtomicCas.Module,
+			features: api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
 			calls: []callCase{
 				// no store
 				{
@@ -811,10 +810,9 @@ func TestE2E(t *testing.T) {
 		{
 			// Checks if load works when comparison value is zero. It wouldn't if
 			// the zero register gets used.
-			name:      "atomic_cas_const0",
-			m:         testcases.AtomicCasConst0.Module,
-			features:  api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
-			skipAMD64: true,
+			name:     "atomic_cas_const0",
+			m:        testcases.AtomicCasConst0.Module,
+			features: api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
 			setupMemory: func(mem api.Memory) {
 				mem.WriteUint32Le(0, 1)
 				mem.WriteUint32Le(8, 2)