wazevo(arm64): support atomic fence and enable tests (#2060)

internal/engine/wazevo/backend/isa/arm64/instr.go

@@ -163,6 +163,7 @@ var defKinds = [numInstructionKinds]defKind{
 	atomicCas:            defKindNone,
 	atomicLoad:           defKindRD,
 	atomicStore:          defKindNone,
+	dmb:                  defKindNone,
 }
 
 // Defs returns the list of regalloc.VReg that are defined by the instruction.
@@ -299,6 +300,7 @@ var useKinds = [numInstructionKinds]useKind{
 	atomicCas:            useKindRDRewrite,
 	atomicLoad:           useKindRN,
 	atomicStore:          useKindRNRM,
+	dmb:                  useKindNone,
 }
 
 // Uses returns the list of regalloc.VReg that are used by the instruction.
@@ -1506,6 +1508,8 @@ func (i *instruction) String() (str string) {
 			m = m + "b"
 		}
 		str = fmt.Sprintf("%s %s, %s", m, formatVRegSized(i.rm.nr(), size), formatVRegSized(i.rn.nr(), 64))
+	case dmb:
+		str = "dmb"
 	case udf:
 		str = "udf"
 	case emitSourceOffsetInfo:
@@ -1549,6 +1553,10 @@ func (i *instruction) asAtomicStore(rn, rt operand, size uint64) {
 	i.u2 = size
 }
 
+func (i *instruction) asDMB() {
+	i.kind = dmb
+}
+
 // TODO: delete unnecessary things.
 const (
 	// nop0 represents a no-op of zero size.
@@ -1724,6 +1732,8 @@ const (
 	atomicLoad
 	// atomicStore represents an atomic store with two source registers and no destination.
 	atomicStore
+	// dmb represents the data memory barrier instruction in inner-shareable (ish) mode.
+	dmb
 	// UDF is the undefined instruction. For debugging only.
 	udf
 

internal/engine/wazevo/backend/isa/arm64/instr_encoding.go

@@ -411,6 +411,8 @@ func (i *instruction) encode(m *machine) {
 			uint32(i.u2),
 			0,
 		))
+	case dmb:
+		c.Emit4Bytes(encodeDMB())
 	default:
 		panic(i.String())
 	}
@@ -2333,3 +2335,7 @@ func encodeAtomicLoadStore(rn, rt, size, l uint32) uint32 {
 
 	return _31to21<<21 | _20to16<<16 | _15to10<<10 | rn<<5 | rt
 }
+
+func encodeDMB() uint32 {
+	return 0b11010101000000110011101110111111
+}

internal/engine/wazevo/backend/isa/arm64/instr_encoding_test.go

@@ -1799,6 +1799,9 @@ func TestInstruction_encode(t *testing.T) {
 		{want: "01fc9f08", setup: func(i *instruction) {
 			i.asAtomicStore(operandNR(x0VReg), operandNR(x1VReg), 1)
 		}},
+		{want: "bf3b03d5", setup: func(i *instruction) {
+			i.asDMB()
+		}},
 		{want: "4201231e4201631e4201239e4201639e4201221e4201621e4201229e4201629e", setup: func(i *instruction) {
 			i.asNop0()
 			cur := i

internal/engine/wazevo/backend/isa/arm64/lower_instr.go

@@ -756,6 +756,11 @@ func (m *machine) LowerInstr(instr *ssa.Instruction) {
 	case ssa.OpcodeAtomicStore:
 		m.lowerAtomicStore(instr)
 
+	case ssa.OpcodeFence:
+		instr := m.allocateInstr()
+		instr.asDMB()
+		m.insert(instr)
+
 	default:
 		panic("TODO: lowering " + op.String())
 	}

internal/engine/wazevo/e2e_test.go

@@ -850,6 +850,15 @@ func TestE2E(t *testing.T) {
 				{params: []uint64{10, 20, 30, 40, 50, 60, 70}, expResults: []uint64{10, 20, 30, 40, 50, 60, 70}},
 			},
 		},
+		{
+			name:      "atomic_fence",
+			m:         testcases.AtomicFence.Module,
+			features:  api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
+			skipAMD64: true,
+			calls: []callCase{
+				{params: []uint64{}, expResults: []uint64{}},
+			},
+		},
 		{
 			name: "float_le",
 			m:    testcases.FloatLe.Module,

internal/engine/wazevo/frontend/frontend_test.go

@@ -2712,6 +2712,7 @@ blk0: (exec_ctx:i64, module_ctx:i64, v2:i32, v3:i32, v4:i32, v5:i32, v6:i32, v7:
 			exp: `
 blk0: (exec_ctx:i64, module_ctx:i64)
 	Fence 0
+	Jump blk_ret
 `,
 		},
 		{
@@ -2720,6 +2721,7 @@ blk0: (exec_ctx:i64, module_ctx:i64)
 			features: api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
 			exp: `
 blk0: (exec_ctx:i64, module_ctx:i64)
+	Jump blk_ret
 `,
 		},
 	} {

internal/engine/wazevo/frontend/lower.go

@@ -3378,7 +3378,6 @@ func (c *Compiler) lowerCurrentOpcode() {
 			state.push(res)
 		case wasm.OpcodeAtomicFence:
 			order := c.readByte()
-			c.readI32u()
 			if state.unreachable {
 				break
 			}

internal/integration_test/engine/threads_test.go

@@ -2,11 +2,13 @@ package adhoc
 
 import (
 	_ "embed"
+	"runtime"
 	"testing"
 
 	wazero "github.com/tetratelabs/wazero"
 	"github.com/tetratelabs/wazero/api"
 	"github.com/tetratelabs/wazero/experimental"
+	"github.com/tetratelabs/wazero/experimental/opt"
 	"github.com/tetratelabs/wazero/internal/platform"
 	"github.com/tetratelabs/wazero/internal/testing/hammer"
 	"github.com/tetratelabs/wazero/internal/testing/require"
@@ -64,6 +66,13 @@ func TestThreadsInterpreter_hammer(t *testing.T) {
 	runAllTests(t, threadTests, wazero.NewRuntimeConfigInterpreter().WithCoreFeatures(api.CoreFeaturesV2|experimental.CoreFeaturesThreads), false)
 }
 
+func TestThreadsWazevo(t *testing.T) {
+	if !platform.CompilerSupported() || runtime.GOARCH != "arm64" {
+		t.Skip()
+	}
+	runAllTests(t, threadTests, opt.NewRuntimeConfigOptimizingCompiler().WithCoreFeatures(api.CoreFeaturesV2|experimental.CoreFeaturesThreads), false)
+}
+
 func incrementGuardedByMutex(t *testing.T, r wazero.Runtime) {
 	P := 8               // max count of goroutines
 	if testing.Short() { // Adjust down if `-test.short`

internal/integration_test/spectest/threads/spec_test.go

@@ -3,11 +3,13 @@ package spectest
 import (
 	"context"
 	"embed"
+	"runtime"
 	"testing"
 
 	"github.com/tetratelabs/wazero"
 	"github.com/tetratelabs/wazero/api"
 	"github.com/tetratelabs/wazero/experimental"
+	"github.com/tetratelabs/wazero/experimental/opt"
 	"github.com/tetratelabs/wazero/internal/integration_test/spectest"
 	"github.com/tetratelabs/wazero/internal/platform"
 )
@@ -28,3 +30,10 @@ func TestCompiler(t *testing.T) {
 func TestInterpreter(t *testing.T) {
 	spectest.Run(t, testcases, context.Background(), wazero.NewRuntimeConfigInterpreter().WithCoreFeatures(enabledFeatures))
 }
+
+func TestWazevo(t *testing.T) {
+	if !platform.CompilerSupported() || runtime.GOARCH != "arm64" {
+		t.Skip()
+	}
+	spectest.Run(t, testcases, context.Background(), opt.NewRuntimeConfigOptimizingCompiler().WithCoreFeatures(enabledFeatures))
+}