Repository: tetratelabs/wazero
Pull request #2113 (2024-03-04 05:02:15)

wazevo(amd64): support for atomic stores
cc @anuraaga 


Patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr.go b/internal/engine/wazevo/backend/isa/amd64/instr.go
index 600ae2b1..74d66f5a 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr.go
@@ -226,7 +226,18 @@ func (i *instruction) String() string {
 	case callIndirect:
 		return fmt.Sprintf("callq *%s", i.op1.format(true))
 	case xchg:
-		return fmt.Sprintf("xchg %s, %s", i.op1.format(true), i.op2.format(true))
+		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("xchg.%s %s, %s", suffix, i.op1.format(true), i.op2.format(true))
 	case zeros:
 		return fmt.Sprintf("xor %s, %s", i.op2.format(true), i.op2.format(true))
 	case fcvtToSintSequence:
@@ -723,9 +734,9 @@ const (
 	// An instruction that will always trigger the illegal instruction exception.
 	ud2
 
-	// xchg swaps the contents of two gp registers.
-	// The instruction doesn't make sense before register allocation, so it doensn't
-	// have useKinds and defKinds to avoid being used by the register allocator.
+	// xchg is described in https://www.felixcloutier.com/x86/xchg.
+	// This instruction uses two operands, where one of them can be a memory address, and swaps their values.
+	// If the dst is a memory address, the execution is atomic.
 	xchg
 
 	// zeros puts zeros into the destination register. This is implemented as xor reg, reg for
@@ -1418,10 +1429,11 @@ func (i *instruction) asPush64(op operand) *instruction {
 	return i
 }
 
-func (i *instruction) asXCHG(rm, rd regalloc.VReg) *instruction {
+func (i *instruction) asXCHG(rm regalloc.VReg, rd operand, size byte) *instruction {
 	i.kind = xchg
 	i.op1 = newOperandReg(rm)
-	i.op2 = newOperandReg(rd)
+	i.op2 = rd
+	i.u1 = uint64(size)
 	return i
 }
 
@@ -2231,6 +2243,7 @@ var defKinds = [instrMax]defKind{
 	idivRemSequence:        defKindDivRem,
 	blendvpd:               defKindNone,
 	mfence:                 defKindNone,
+	xchg:                   defKindNone,
 }
 
 // String implements fmt.Stringer.
@@ -2307,6 +2320,7 @@ var useKinds = [instrMax]useKind{
 	idivRemSequence:        useKindDivRem,
 	blendvpd:               useKindBlendvpd,
 	mfence:                 useKindNone,
+	xchg:                   useKindOp1RegOp2,
 }
 
 func (u useKind) String() string {
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go b/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
index 2f99e8b5..201f18de 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
@@ -1245,10 +1245,46 @@ func (i *instruction) encode(c backend.Compiler) (needsLabelResolution bool) {
 		}
 
 	case xchg:
-		r1, r2 := i.op1.reg().RealReg(), i.op2.reg().RealReg()
+		src, dst := regEncodings[i.op1.reg().RealReg()], i.op2
+		size := i.u1
 
-		enc1, enc2 := regEncodings[r1], regEncodings[r2]
-		encodeRegReg(c, legacyPrefixesNone, 0x87, 1, enc2, enc1, rexInfo(0).setW())
+		var rex rexInfo
+		var opcode uint32
+		lp := legacyPrefixesNone
+		switch size {
+		case 8:
+			opcode = 0x87
+			rex = rexInfo(0).setW()
+		case 4:
+			opcode = 0x87
+			rex = rexInfo(0).clearW()
+		case 2:
+			lp = legacyPrefixes0x66
+			opcode = 0x87
+			rex = rexInfo(0).clearW()
+		case 1:
+			opcode = 0x86
+			if i.op2.kind == operandKindReg {
+				panic("TODO?: xchg on two 1-byte registers")
+			}
+			// 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, 1, src, m, rex)
+		case operandKindReg:
+			r := dst.reg().RealReg()
+			encodeRegReg(c, lp, opcode, 1, src, regEncodings[r], rex)
+		default:
+			panic("BUG: invalid operand kind")
+		}
 
 	case zeros:
 		r := i.op2.reg()
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine.go b/internal/engine/wazevo/backend/isa/amd64/machine.go
index 29575ee0..f962526b 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine.go
@@ -959,6 +959,18 @@ func (m *machine) LowerInstr(instr *ssa.Instruction) {
 	case ssa.OpcodeFence:
 		m.insert(m.allocateInstr().asMFence())
 
+	case ssa.OpcodeAtomicStore:
+		ptr, _val := instr.Arg2()
+		size := instr.AtomicTargetSize()
+
+		val := m.getOperand_Reg(m.c.ValueDefinition(_val))
+		// The content on the val register will be overwritten by xchg, so we need to copy it to a temporary register.
+		copied := m.copyToTmp(val.reg())
+
+		mem := newOperandMem(m.lowerToAddressMode(ptr, 0))
+		store := m.allocateInstr().asXCHG(copied, mem, byte(size))
+		m.insert(store)
+
 	default:
 		panic("TODO: lowering " + op.String())
 	}
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine_regalloc.go b/internal/engine/wazevo/backend/isa/amd64/machine_regalloc.go
index e2d01362..0bb28ee9 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine_regalloc.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine_regalloc.go
@@ -107,7 +107,7 @@ func (m *machine) ClobberedRegisters(regs []regalloc.VReg) {
 func (m *machine) Swap(cur *instruction, x1, x2, tmp regalloc.VReg) {
 	if x1.RegType() == regalloc.RegTypeInt {
 		prevNext := cur.next
-		xc := m.allocateInstr().asXCHG(x1, x2)
+		xc := m.allocateInstr().asXCHG(x1, newOperandReg(x2), 8)
 		cur = linkInstr(cur, xc)
 		linkInstr(cur, prevNext)
 	} else {
diff --git a/internal/engine/wazevo/ssa/pass.go b/internal/engine/wazevo/ssa/pass.go
index 69afa7cb..a2e986cd 100644
--- a/internal/engine/wazevo/ssa/pass.go
+++ b/internal/engine/wazevo/ssa/pass.go
@@ -112,88 +112,105 @@ func passDeadBlockEliminationOpt(b *builder) {
 func passRedundantPhiEliminationOpt(b *builder) {
 	redundantParameterIndexes := b.ints[:0] // reuse the slice from previous iterations.
 
-	_ = b.blockIteratorBegin() // skip entry block!
-	// Below, we intentionally use the named iteration variable name, as this comes with inevitable nested for loops!
-	for blk := b.blockIteratorNext(); blk != nil; blk = b.blockIteratorNext() {
-		paramNum := len(blk.params)
+	// TODO: this might be costly for large programs, but at least, as far as I did the experiment, it's almost the
+	//  same as the single iteration version in terms of the overall compilation time. That *might be* mostly thanks to the fact
+	//  that removing many PHIs results in the reduction of the total instructions, not because of this indefinite iteration is
+	//  relatively small. For example, sqlite speedtest binary results in the large number of redundant PHIs,
+	//  the maximum number of iteration was 22, which seems to be acceptable but not that small either since the
+	//  complexity here is O(BlockNum * Iterations) at the worst case where BlockNum might be the order of thousands.
+	for {
+		changed := false
+		_ = b.blockIteratorBegin() // skip entry block!
+		// Below, we intentionally use the named iteration variable name, as this comes with inevitable nested for loops!
+		for blk := b.blockIteratorNext(); blk != nil; blk = b.blockIteratorNext() {
+			paramNum := len(blk.params)
+
+			for paramIndex := 0; paramIndex < paramNum; paramIndex++ {
+				phiValue := blk.params[paramIndex].value
+				redundant := true
+
+				nonSelfReferencingValue := ValueInvalid
+				for predIndex := range blk.preds {
+					br := blk.preds[predIndex].branch
+					// Resolve the alias in the arguments so that we could use the previous iteration's result.
+					b.resolveArgumentAlias(br)
+					pred := br.vs.View()[paramIndex]
+					if pred == phiValue {
+						// This is self-referencing: PHI from the same PHI.
+						continue
+					}
 
-		for paramIndex := 0; paramIndex < paramNum; paramIndex++ {
-			phiValue := blk.params[paramIndex].value
-			redundant := true
+					if !nonSelfReferencingValue.Valid() {
+						nonSelfReferencingValue = pred
+						continue
+					}
 
-			nonSelfReferencingValue := ValueInvalid
-			for predIndex := range blk.preds {
-				pred := blk.preds[predIndex].branch.vs.View()[paramIndex]
-				if pred == phiValue {
-					// This is self-referencing: PHI from the same PHI.
-					continue
+					if nonSelfReferencingValue != pred {
+						redundant = false
+						break
+					}
 				}
 
 				if !nonSelfReferencingValue.Valid() {
-					nonSelfReferencingValue = pred
-					continue
+					// This shouldn't happen, and must be a bug in builder.go.
+					panic("BUG: params added but only self-referencing")
 				}
 
-				if nonSelfReferencingValue != pred {
-					redundant = false
-					break
+				if redundant {
+					b.redundantParameterIndexToValue[paramIndex] = nonSelfReferencingValue
+					redundantParameterIndexes = append(redundantParameterIndexes, paramIndex)
 				}
 			}
 
-			if !nonSelfReferencingValue.Valid() {
-				// This shouldn't happen, and must be a bug in builder.go.
-				panic("BUG: params added but only self-referencing")
+			if len(b.redundantParameterIndexToValue) == 0 {
+				continue
 			}
+			changed = true
 
-			if redundant {
-				b.redundantParameterIndexToValue[paramIndex] = nonSelfReferencingValue
-				redundantParameterIndexes = append(redundantParameterIndexes, paramIndex)
+			// Remove the redundant PHIs from the argument list of branching instructions.
+			for predIndex := range blk.preds {
+				var cur int
+				predBlk := blk.preds[predIndex]
+				branchInst := predBlk.branch
+				view := branchInst.vs.View()
+				for argIndex, value := range view {
+					if _, ok := b.redundantParameterIndexToValue[argIndex]; !ok {
+						view[cur] = value
+						cur++
+					}
+				}
+				branchInst.vs.Cut(cur)
 			}
-		}
 
-		if len(b.redundantParameterIndexToValue) == 0 {
-			continue
-		}
+			// Still need to have the definition of the value of the PHI (previously as the parameter).
+			for _, redundantParamIndex := range redundantParameterIndexes {
+				phiValue := blk.params[redundantParamIndex].value
+				onlyValue := b.redundantParameterIndexToValue[redundantParamIndex]
+				// Create an alias in this block from the only phi argument to the phi value.
+				b.alias(phiValue, onlyValue)
+			}
 
-		// Remove the redundant PHIs from the argument list of branching instructions.
-		for predIndex := range blk.preds {
+			// Finally, Remove the param from the blk.
 			var cur int
-			predBlk := blk.preds[predIndex]
-			branchInst := predBlk.branch
-			view := branchInst.vs.View()
-			for argIndex, value := range view {
-				if _, ok := b.redundantParameterIndexToValue[argIndex]; !ok {
-					view[cur] = value
+			for paramIndex := 0; paramIndex < paramNum; paramIndex++ {
+				param := blk.params[paramIndex]
+				if _, ok := b.redundantParameterIndexToValue[paramIndex]; !ok {
+					blk.params[cur] = param
 					cur++
 				}
 			}
-			branchInst.vs.Cut(cur)
-		}
-
-		// Still need to have the definition of the value of the PHI (previously as the parameter).
-		for _, redundantParamIndex := range redundantParameterIndexes {
-			phiValue := blk.params[redundantParamIndex].value
-			onlyValue := b.redundantParameterIndexToValue[redundantParamIndex]
-			// Create an alias in this block from the only phi argument to the phi value.
-			b.alias(phiValue, onlyValue)
-		}
+			blk.params = blk.params[:cur]
 
-		// Finally, Remove the param from the blk.
-		var cur int
-		for paramIndex := 0; paramIndex < paramNum; paramIndex++ {
-			param := blk.params[paramIndex]
-			if _, ok := b.redundantParameterIndexToValue[paramIndex]; !ok {
-				blk.params[cur] = param
-				cur++
+			// Clears the map for the next iteration.
+			for _, paramIndex := range redundantParameterIndexes {
+				delete(b.redundantParameterIndexToValue, paramIndex)
 			}
+			redundantParameterIndexes = redundantParameterIndexes[:0]
 		}
-		blk.params = blk.params[:cur]
 
-		// Clears the map for the next iteration.
-		for _, paramIndex := range redundantParameterIndexes {
-			delete(b.redundantParameterIndexToValue, paramIndex)
+		if !changed {
+			break
 		}
-		redundantParameterIndexes = redundantParameterIndexes[:0]
 	}
 
 	// Reuse the slice for the future passes.


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 be85fa7e..2d1b7f13 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr_encoding_test.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr_encoding_test.go
@@ -3895,19 +3895,89 @@ func TestInstruction_format_encode(t *testing.T) {
 			wantFormat: "exit_sequence %r15",
 		},
 		{
-			setup:      func(i *instruction) { i.asXCHG(r11VReg, r14VReg) },
-			want:       "4d87f3",
-			wantFormat: "xchg %r11, %r14",
+			setup:      func(i *instruction) { i.asXCHG(r11VReg, newOperandReg(r14VReg), 8) },
+			want:       "4d87de",
+			wantFormat: "xchg.q %r11, %r14",
 		},
 		{
-			setup:      func(i *instruction) { i.asXCHG(r15VReg, raxVReg) },
-			want:       "4987c7",
-			wantFormat: "xchg %r15, %rax",
+			setup:      func(i *instruction) { i.asXCHG(r11VReg, newOperandMem(newAmodeImmReg(123, raxVReg)), 8) },
+			want:       "4c87587b",
+			wantFormat: "xchg.q %r11, 123(%rax)",
 		},
 		{
-			setup:      func(i *instruction) { i.asXCHG(rbxVReg, rsiVReg) },
-			want:       "4887f3",
-			wantFormat: "xchg %rbx, %rsi",
+			setup:      func(i *instruction) { i.asXCHG(r15VReg, newOperandReg(raxVReg), 8) },
+			want:       "4c87f8",
+			wantFormat: "xchg.q %r15, %rax",
+		},
+		{
+			setup:      func(i *instruction) { i.asXCHG(rbxVReg, newOperandReg(rsiVReg), 8) },
+			want:       "4887de",
+			wantFormat: "xchg.q %rbx, %rsi",
+		},
+		{
+			setup:      func(i *instruction) { i.asXCHG(r11VReg, newOperandReg(r14VReg), 4) },
+			want:       "4587de",
+			wantFormat: "xchg.l %r11, %r14",
+		},
+		{
+			setup:      func(i *instruction) { i.asXCHG(r15VReg, newOperandReg(raxVReg), 4) },
+			want:       "4487f8",
+			wantFormat: "xchg.l %r15, %rax",
+		},
+		{
+			setup:      func(i *instruction) { i.asXCHG(rbxVReg, newOperandReg(rsiVReg), 4) },
+			want:       "87de",
+			wantFormat: "xchg.l %rbx, %rsi",
+		},
+		{
+			setup:      func(i *instruction) { i.asXCHG(r11VReg, newOperandMem(newAmodeImmReg(123, raxVReg)), 4) },
+			want:       "4487587b",
+			wantFormat: "xchg.l %r11, 123(%rax)",
+		},
+		{
+			setup:      func(i *instruction) { i.asXCHG(r11VReg, newOperandReg(r14VReg), 2) },
+			want:       "664587de",
+			wantFormat: "xchg.w %r11, %r14",
+		},
+		{
+			setup:      func(i *instruction) { i.asXCHG(r15VReg, newOperandReg(raxVReg), 2) },
+			want:       "664487f8",
+			wantFormat: "xchg.w %r15, %rax",
+		},
+		{
+			setup:      func(i *instruction) { i.asXCHG(rbxVReg, newOperandReg(rsiVReg), 2) },
+			want:       "6687de",
+			wantFormat: "xchg.w %rbx, %rsi",
+		},
+		{
+			setup:      func(i *instruction) { i.asXCHG(r11VReg, newOperandMem(newAmodeImmReg(123, raxVReg)), 2) },
+			want:       "664487587b",
+			wantFormat: "xchg.w %r11, 123(%rax)",
+		},
+		{
+			setup:      func(i *instruction) { i.asXCHG(r11VReg, newOperandMem(newAmodeImmReg(123, raxVReg)), 1) },
+			want:       "4486587b",
+			wantFormat: "xchg.b %r11, 123(%rax)",
+		},
+		{
+			setup:      func(i *instruction) { i.asXCHG(r11VReg, newOperandMem(newAmodeImmReg(123, rdiVReg)), 1) },
+			want:       "44865f7b",
+			wantFormat: "xchg.b %r11, 123(%rdi)",
+		},
+		{
+			setup:      func(i *instruction) { i.asXCHG(rsiVReg, newOperandMem(newAmodeImmReg(123, raxVReg)), 1) },
+			want:       "4086707b",
+			wantFormat: "xchg.b %rsi, 123(%rax)",
+		},
+		{
+			setup:      func(i *instruction) { i.asXCHG(rdiVReg, newOperandMem(newAmodeImmReg(123, raxVReg)), 1) },
+			want:       "4086787b",
+			wantFormat: "xchg.b %rdi, 123(%rax)",
+		},
+		{
+			setup:      func(i *instruction) { i.asXCHG(rdiVReg, newOperandMem(newAmodeImmReg(123, rsiVReg)), 1) },
+			want:       "40867e7b",
+			wantFormat: "xchg.b %rdi, 123(%rsi)",
 		},
 		{
 			setup:      func(i *instruction) { i.asZeros(rbxVReg) },
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine_regalloc_test.go b/internal/engine/wazevo/backend/isa/amd64/machine_regalloc_test.go
index 3eeb99e1..929e14c3 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine_regalloc_test.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine_regalloc_test.go
@@ -172,7 +172,7 @@ func TestMachineSwap(t *testing.T) {
 			tmp: rdiVReg,
 			expected: `
 	ud2
-	xchg %r15, %rax
+	xchg.q %r15, %rax
 	exit_sequence %rsi
 `,
 		},
@@ -181,7 +181,7 @@ func TestMachineSwap(t *testing.T) {
 			x2: raxVReg,
 			expected: `
 	ud2
-	xchg %r15, %rax
+	xchg.q %r15, %rax
 	exit_sequence %rsi
 `,
 		},
diff --git a/internal/engine/wazevo/e2e_test.go b/internal/engine/wazevo/e2e_test.go
index 9807ebbe..1adeabb9 100644
--- a/internal/engine/wazevo/e2e_test.go
+++ b/internal/engine/wazevo/e2e_test.go
@@ -832,10 +832,9 @@ func TestE2E(t *testing.T) {
 			},
 		},
 		{
-			name:      "atomic_store_load",
-			m:         testcases.AtomicStoreLoad.Module,
-			features:  api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
-			skipAMD64: true,
+			name:     "atomic_store_load",
+			m:        testcases.AtomicStoreLoad.Module,
+			features: api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
 			calls: []callCase{
 				{params: []uint64{1, 2, 3, 4, 5, 6, 7}, expResults: []uint64{1, 2, 3, 4, 5, 6, 7}},
 				{params: []uint64{10, 20, 30, 40, 50, 60, 70}, expResults: []uint64{10, 20, 30, 40, 50, 60, 70}},
diff --git a/internal/engine/wazevo/ssa/pass_test.go b/internal/engine/wazevo/ssa/pass_test.go
index 261f3225..83d33fbe 100644
--- a/internal/engine/wazevo/ssa/pass_test.go
+++ b/internal/engine/wazevo/ssa/pass_test.go
@@ -189,7 +189,7 @@ blk0: ()
 
 blk1: (v0:i32) <-- (blk0,blk1)
 	v3:i32 = Iconst_32 0xff
-	Brz v2, blk1, v3
+	Brz v1, blk1, v3
 	Jump blk2
 
 blk2: () <-- (blk1)