Repository: tetratelabs/wazero
Pull request #2124 (2024-03-06 05:46:43)

wazevo(amd64): complete threads spec
😎 


Patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr.go b/internal/engine/wazevo/backend/isa/amd64/instr.go
index f0d272eb..137994c0 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr.go
@@ -297,6 +297,10 @@ func (i *instruction) String() string {
 			suffix = "q"
 		}
 		return fmt.Sprintf("lock xadd.%s %s, %s", suffix, i.op1.format(true), i.op2.format(true))
+
+	case keepAlive:
+		return fmt.Sprintf("keepAlive %s", i.op1.format(true))
+
 	default:
 		panic(fmt.Sprintf("BUG: %d", int(i.kind)))
 	}
@@ -856,6 +860,9 @@ const (
 	// lockxadd is xadd https://www.felixcloutier.com/x86/xadd with a lock prefix.
 	lockxadd
 
+	// keepAlive is a meta instruction that uses one register and does nothing.
+	keepAlive
+
 	instrMax
 )
 
@@ -864,6 +871,12 @@ func (i *instruction) asMFence() *instruction {
 	return i
 }
 
+func (i *instruction) asKeepAlive(r regalloc.VReg) *instruction {
+	i.kind = keepAlive
+	i.op1 = newOperandReg(r)
+	return i
+}
+
 func (i *instruction) asIdivRemSequence(execCtx, divisor, tmpGp regalloc.VReg, isDiv, signed, _64 bool) *instruction {
 	i.kind = idivRemSequence
 	i.op1 = newOperandReg(execCtx)
@@ -2353,6 +2366,7 @@ var defKinds = [instrMax]defKind{
 	lockcmpxchg:            defKindNone,
 	lockxadd:               defKindNone,
 	neg:                    defKindNone,
+	keepAlive:              defKindNone,
 }
 
 // String implements fmt.Stringer.
@@ -2435,6 +2449,7 @@ var useKinds = [instrMax]useKind{
 	lockcmpxchg:            useKindRaxOp1RegOp2,
 	lockxadd:               useKindOp1RegOp2,
 	neg:                    useKindOp1,
+	keepAlive:              useKindOp1,
 }
 
 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 80f766bf..1aabaa8c 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
@@ -11,7 +11,7 @@ import (
 
 func (i *instruction) encode(c backend.Compiler) (needsLabelResolution bool) {
 	switch kind := i.kind; kind {
-	case nop0, sourceOffsetInfo, defineUninitializedReg, fcvtToSintSequence, fcvtToUintSequence:
+	case nop0, sourceOffsetInfo, defineUninitializedReg, fcvtToSintSequence, fcvtToUintSequence, keepAlive:
 	case ret:
 		encodeRet(c)
 	case imm:
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine.go b/internal/engine/wazevo/backend/isa/amd64/machine.go
index 1f5728d4..3165c397 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine.go
@@ -989,40 +989,79 @@ func (m *machine) LowerInstr(instr *ssa.Instruction) {
 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:
+		valCopied := m.copyToTmp(_val.reg())
 		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.clearHigherBitsForAtomic(valCopied, size, ret.Type())
 		m.copyTo(valCopied, m.c.VRegOf(ret))
 
-	case ssa.AtomicRmwOpAnd:
-		panic("TODO")
+	case ssa.AtomicRmwOpAnd, ssa.AtomicRmwOpOr, ssa.AtomicRmwOpXor:
+		accumulator := raxVReg
+		// Reserve rax for the accumulator to make regalloc happy.
+		// Note: do this initialization before defining valCopied, because it might be the same register and
+		// if that happens, the unnecessary load/store will be performed inside the loop.
+		// This can be mitigated in any way once the register allocator is clever enough.
+		m.insert(m.allocateInstr().asDefineUninitializedReg(accumulator))
+
+		// Copy the value to a temporary register.
+		valCopied := m.copyToTmp(_val.reg())
+		m.clearHigherBitsForAtomic(valCopied, size, ret.Type())
+
+		memOp := newOperandMem(mem)
+		tmp := m.c.AllocateVReg(ssa.TypeI64)
+		beginLoop, beginLoopLabel := m.allocateBrTarget()
+		{
+			m.insert(beginLoop)
+			// Reset the value on tmp by the original value.
+			m.copyTo(valCopied, tmp)
+			// Load the current value at the memory location into accumulator.
+			switch size {
+			case 1:
+				m.insert(m.allocateInstr().asMovzxRmR(extModeBQ, memOp, accumulator))
+			case 2:
+				m.insert(m.allocateInstr().asMovzxRmR(extModeWQ, memOp, accumulator))
+			case 4:
+				m.insert(m.allocateInstr().asMovzxRmR(extModeLQ, memOp, accumulator))
+			case 8:
+				m.insert(m.allocateInstr().asMov64MR(memOp, accumulator))
+			default:
+				panic("BUG")
+			}
+			// Then perform the logical operation on the accumulator and the value on tmp.
+			switch op {
+			case ssa.AtomicRmwOpAnd:
+				m.insert(m.allocateInstr().asAluRmiR(aluRmiROpcodeAnd, newOperandReg(accumulator), tmp, true))
+			case ssa.AtomicRmwOpOr:
+				m.insert(m.allocateInstr().asAluRmiR(aluRmiROpcodeOr, newOperandReg(accumulator), tmp, true))
+			case ssa.AtomicRmwOpXor:
+				m.insert(m.allocateInstr().asAluRmiR(aluRmiROpcodeXor, newOperandReg(accumulator), tmp, true))
+			default:
+				panic("BUG")
+			}
+			// Finally, try compare-exchange the value at the memory location with the tmp.
+			m.insert(m.allocateInstr().asLockCmpXCHG(tmp, memOp.addressMode(), byte(size)))
+			// If it succeeds, ZF will be set, and we can break the loop.
+			m.insert(m.allocateInstr().asJmpIf(condNZ, newOperandLabel(beginLoopLabel)))
+		}
 
-	case ssa.AtomicRmwOpOr:
-		panic("TODO")
+		// valCopied must be alive at the end of the loop.
+		m.insert(m.allocateInstr().asKeepAlive(valCopied))
 
-	case ssa.AtomicRmwOpXor:
-		panic("TODO")
+		// At this point, accumulator contains the result.
+		m.clearHigherBitsForAtomic(accumulator, size, ret.Type())
+		m.copyTo(accumulator, m.c.VRegOf(ret))
 
 	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))
-		}
+		valCopied := m.copyToTmp(_val.reg())
 
+		m.insert(m.allocateInstr().asXCHG(valCopied, newOperandMem(mem), byte(size)))
+		m.clearHigherBitsForAtomic(valCopied, size, ret.Type())
 		m.copyTo(valCopied, m.c.VRegOf(ret))
 
 	default:
@@ -1038,13 +1077,29 @@ func (m *machine) lowerAtomicCas(addr, exp, repl ssa.Value, size uint64, ret ssa
 	accumulator := raxVReg
 	m.copyTo(expOp.reg(), accumulator)
 	m.insert(m.allocateInstr().asLockCmpXCHG(replOp.reg(), mem, byte(size)))
+	m.clearHigherBitsForAtomic(accumulator, size, ret.Type())
+	m.copyTo(accumulator, m.c.VRegOf(ret))
+}
 
-	if size < 4 {
-		// Clear the unnecessary bits.
-		m.insert(m.allocateInstr().asAluRmiR(aluRmiROpcodeAnd, newOperandImm32(uint32((1<<(8*size))-1)), accumulator, true))
+func (m *machine) clearHigherBitsForAtomic(r regalloc.VReg, valSize uint64, resultType ssa.Type) {
+	switch resultType {
+	case ssa.TypeI32:
+		switch valSize {
+		case 1:
+			m.insert(m.allocateInstr().asMovzxRmR(extModeBL, newOperandReg(r), r))
+		case 2:
+			m.insert(m.allocateInstr().asMovzxRmR(extModeWL, newOperandReg(r), r))
+		}
+	case ssa.TypeI64:
+		switch valSize {
+		case 1:
+			m.insert(m.allocateInstr().asMovzxRmR(extModeBQ, newOperandReg(r), r))
+		case 2:
+			m.insert(m.allocateInstr().asMovzxRmR(extModeWQ, newOperandReg(r), r))
+		case 4:
+			m.insert(m.allocateInstr().asMovzxRmR(extModeLQ, newOperandReg(r), r))
+		}
 	}
-
-	m.copyTo(accumulator, m.c.VRegOf(ret))
 }
 
 func (m *machine) lowerFcmp(instr *ssa.Instruction) {


Test patch:
diff --git a/internal/engine/wazevo/e2e_test.go b/internal/engine/wazevo/e2e_test.go
index 21dadb86..6536f915 100644
--- a/internal/engine/wazevo/e2e_test.go
+++ b/internal/engine/wazevo/e2e_test.go
@@ -6,7 +6,6 @@ import (
 	"encoding/binary"
 	"fmt"
 	"math"
-	"runtime"
 	"testing"
 
 	"github.com/tetratelabs/wazero"
@@ -30,13 +29,6 @@ const (
 	v128 = wasm.ValueTypeV128
 )
 
-// TODO: delete once we complete the implementation for amd64.
-func skipOnAmd64(t *testing.T) {
-	if runtime.GOARCH == "amd64" {
-		t.Skip("skip on amd64")
-	}
-}
-
 func TestE2E(t *testing.T) {
 	tmp := t.TempDir()
 	type callCase struct {
@@ -49,7 +41,6 @@ func TestE2E(t *testing.T) {
 		imported, m *wasm.Module
 		calls       []callCase
 		features    api.CoreFeatures
-		skipAMD64   bool
 		setupMemory func(mem api.Memory)
 	}{
 		{
@@ -686,10 +677,9 @@ func TestE2E(t *testing.T) {
 			},
 		},
 		{
-			name:      "atomic_rmw_and",
-			m:         testcases.AtomicRmwAnd.Module,
-			features:  api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
-			skipAMD64: true,
+			name:     "atomic_rmw_and",
+			m:        testcases.AtomicRmwAnd.Module,
+			features: api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
 			setupMemory: func(mem api.Memory) {
 				mem.WriteUint32Le(0, 0xffffffff)
 				mem.WriteUint32Le(8, 0xffffffff)
@@ -715,10 +705,9 @@ func TestE2E(t *testing.T) {
 			},
 		},
 		{
-			name:      "atomic_rmw_or",
-			m:         testcases.AtomicRmwOr.Module,
-			features:  api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
-			skipAMD64: true,
+			name:     "atomic_rmw_or",
+			m:        testcases.AtomicRmwOr.Module,
+			features: api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
 			calls: []callCase{
 				{
 					params:     []uint64{0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f},
@@ -735,10 +724,9 @@ func TestE2E(t *testing.T) {
 			},
 		},
 		{
-			name:      "atomic_rmw_xor",
-			m:         testcases.AtomicRmwXor.Module,
-			features:  api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
-			skipAMD64: true,
+			name:     "atomic_rmw_xor",
+			m:        testcases.AtomicRmwXor.Module,
+			features: api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
 			calls: []callCase{
 				{
 					params:     []uint64{0, 0, 0, 0, 0, 0, 0},
@@ -856,9 +844,6 @@ func TestE2E(t *testing.T) {
 	} {
 		tc := tc
 		t.Run(tc.name, func(t *testing.T) {
-			if tc.skipAMD64 {
-				skipOnAmd64(t)
-			}
 			for i := 0; i < 2; i++ {
 				var name string
 				if i == 0 {
diff --git a/internal/engine/wazevo/wazevo_test.go b/internal/engine/wazevo/wazevo_test.go
index 7a59a4b6..4c49ac26 100644
--- a/internal/engine/wazevo/wazevo_test.go
+++ b/internal/engine/wazevo/wazevo_test.go
@@ -3,12 +3,12 @@ package wazevo
 import (
 	"context"
 	"os"
-	"runtime"
 	"testing"
 	"unsafe"
 
 	"github.com/tetratelabs/wazero/api"
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/wazevoapi"
+	"github.com/tetratelabs/wazero/internal/platform"
 	"github.com/tetratelabs/wazero/internal/testing/require"
 	"github.com/tetratelabs/wazero/internal/wasm"
 )
@@ -16,7 +16,7 @@ import (
 var ctx = context.Background()
 
 func TestMain(m *testing.M) {
-	if runtime.GOARCH != "arm64" && runtime.GOARCH != "amd64" {
+	if !platform.CompilerSupported() {
 		os.Exit(0)
 	}
 	os.Exit(m.Run())
diff --git a/internal/integration_test/engine/threads_test.go b/internal/integration_test/engine/threads_test.go
index ee58af31..c6a2a44c 100644
--- a/internal/integration_test/engine/threads_test.go
+++ b/internal/integration_test/engine/threads_test.go
@@ -2,7 +2,6 @@ package adhoc
 
 import (
 	_ "embed"
-	"runtime"
 	"testing"
 
 	wazero "github.com/tetratelabs/wazero"
@@ -67,7 +66,7 @@ func TestThreadsInterpreter_hammer(t *testing.T) {
 }
 
 func TestThreadsWazevo(t *testing.T) {
-	if !platform.CompilerSupported() || runtime.GOARCH != "arm64" {
+	if !platform.CompilerSupported() {
 		t.Skip()
 	}
 	runAllTests(t, threadTests, opt.NewRuntimeConfigOptimizingCompiler().WithCoreFeatures(api.CoreFeaturesV2|experimental.CoreFeaturesThreads), false)
diff --git a/internal/integration_test/spectest/threads/spec_test.go b/internal/integration_test/spectest/threads/spec_test.go
index 6e2b39ba..b0744e5b 100644
--- a/internal/integration_test/spectest/threads/spec_test.go
+++ b/internal/integration_test/spectest/threads/spec_test.go
@@ -3,7 +3,6 @@ package spectest
 import (
 	"context"
 	"embed"
-	"runtime"
 	"testing"
 
 	"github.com/tetratelabs/wazero"
@@ -32,7 +31,7 @@ func TestInterpreter(t *testing.T) {
 }
 
 func TestWazevo(t *testing.T) {
-	if !platform.CompilerSupported() || runtime.GOARCH != "arm64" {
+	if !platform.CompilerSupported() {
 		t.Skip()
 	}
 	spectest.Run(t, testcases, context.Background(), opt.NewRuntimeConfigOptimizingCompiler().WithCoreFeatures(enabledFeatures))