wazevo: pass address.wast (#1649)

Signed-off-by: Takeshi Yoneda <[email protected]>

internal/engine/wazevo/backend/compiler.go

@@ -143,6 +143,9 @@ type compiler struct {
 // Compile implements Compiler.Compile.
 func (c *compiler) Compile() ([]byte, []RelocationInfo, int, error) {
 	c.Lower()
+	if false {
+		fmt.Printf("[[[after lowering]]]%s", c.Format())
+	}
 	c.RegAlloc()
 	c.Finalize()
 	goPreambleSize := c.Encode()

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

@@ -54,6 +54,7 @@ var defKinds = [numInstructionKinds]defKind{
 	aluRRRExtend:    defKindRD,
 	bitRR:           defKindRD,
 	movZ:            defKindRD,
+	movK:            defKindRD,
 	movN:            defKindRD,
 	mov32:           defKindRD,
 	mov64:           defKindRD,
@@ -152,6 +153,7 @@ var useKinds = [numInstructionKinds]useKind{
 	aluRRRExtend:    useKindRNRM,
 	bitRR:           useKindRN,
 	movZ:            useKindNone,
+	movK:            useKindNone,
 	movN:            useKindNone,
 	mov32:           useKindRN,
 	mov64:           useKindRN,

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

@@ -645,7 +645,7 @@ func bitmaskImmediate(c uint64, is64bit bool) (immr, imms, N byte) {
 	}
 
 	var mode byte = 32
-	if is64bit {
+	if is64bit && size == 64 {
 		N, mode = 0b1, 64
 	}
 

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

@@ -226,6 +226,7 @@ func TestInstruction_encode(t *testing.T) {
 		{want: "4f7a40d2", setup: func(i *instruction) { i.asALUBitmaskImm(aluOpEor, x18VReg, x15VReg, 0x7fffffff, true) }},
 		{want: "4f767fd2", setup: func(i *instruction) { i.asALUBitmaskImm(aluOpEor, x18VReg, x15VReg, 0x7ffffffe, true) }},
 		{want: "4fba7fd2", setup: func(i *instruction) { i.asALUBitmaskImm(aluOpEor, x18VReg, x15VReg, 0xfffffffffffe, true) }},
+		{want: "f20300b2", setup: func(i *instruction) { i.asALUBitmaskImm(aluOpOrr, xzrVReg, x18VReg, 0x100000001, true) }},
 		{want: "f21fbf0e", setup: func(i *instruction) { i.asFpuMov64(v18VReg, v31VReg) }},
 		{want: "f21fbf4e", setup: func(i *instruction) { i.asFpuMov128(v18VReg, v31VReg) }},
 		{want: "4000140b", setup: func(i *instruction) {

internal/engine/wazevo/backend/regalloc/regalloc.go

@@ -188,20 +188,13 @@ func (a *Allocator) livenessAnalysis(f Function) {
 				if def.IsRealReg() {
 					info.realRegDefs[def] = append(info.realRegDefs[def], pos)
 				} else {
-					if _, ok := info.defs[def]; ok {
-						// This must be a bug in lowering logic in the ISA-specific code.
-						// In short, a virtual register must be defined only once because that's the invariant of
-						// liveness analysis result used in the current implementation. If you reach here, you can either
-						// use "temporary" physical register (like x27 in AArch64), or simply allocate a new virtual
-						// register for the additional definition. In any ways, such a new virtual register doesn't cost
-						// at all in the final code because they are instantly killed and won't interfere with other
-						// virtual registers later on.
-						//
-						// TODO: this should be enforced globally, not just per block. But that can be costly to check.
-						panic(fmt.Sprintf("BUG: multiple definitions found for a virtual register %s", def.String()))
+					if _, ok := info.defs[def]; !ok {
+						// This means that this VReg is defined multiple times in a series of instructions
+						// e.g. loading arbitrary constant in arm64, and we only need the earliest
+						// definition to construct live range.
+						info.defs[def] = pos
 					}
 
-					info.defs[def] = pos
 					a.vs = append(a.vs, def)
 				}
 			}
@@ -453,7 +446,7 @@ func (a *Allocator) allocateBlockInfo(blockID int) {
 		a.blockInfos = append(a.blockInfos, make([]blockInfo, blockID+1)...)
 	}
 	info := &a.blockInfos[blockID]
-	info.init()
+	a.initBlockInfo(info)
 }
 
 func (a *Allocator) blockInfoAt(blockID int) (info *blockInfo) {
@@ -491,35 +484,52 @@ func (a *Allocator) allocateNode() (n *node) {
 	return
 }
 
-func (i *blockInfo) init() {
+func resetMap[T any](a *Allocator, m map[VReg]T) {
+	a.vs = a.vs[:0]
+	for v := range m {
+		a.vs = append(a.vs, v)
+	}
+	for _, v := range a.vs {
+		delete(m, v)
+	}
+}
+
+func (a *Allocator) initBlockInfo(i *blockInfo) {
 	i.liveNodes = i.liveNodes[:0]
 	if i.liveOuts == nil {
-		// TODO: reuse!!
 		i.liveOuts = make(map[VReg]struct{})
+	} else {
+		resetMap(a, i.liveOuts)
 	}
 	if i.liveIns == nil {
-		// TODO: reuse!!
 		i.liveIns = make(map[VReg]struct{})
+	} else {
+		resetMap(a, i.liveIns)
 	}
 	if i.defs == nil {
-		// TODO: reuse!!
 		i.defs = make(map[VReg]programCounter)
+	} else {
+		resetMap(a, i.defs)
 	}
 	if i.lastUses == nil {
-		// TODO: reuse!!
 		i.lastUses = make(map[VReg]programCounter)
+	} else {
+		resetMap(a, i.lastUses)
 	}
 	if i.kills == nil {
-		// TODO: reuse!!
 		i.kills = make(map[VReg]programCounter)
+	} else {
+		resetMap(a, i.kills)
 	}
 	if i.realRegUses == nil {
-		// TODO: reuse!!
 		i.realRegUses = make(map[VReg][]programCounter)
+	} else {
+		resetMap(a, i.realRegUses)
 	}
 	if i.realRegDefs == nil {
-		// TODO: reuse!!
 		i.realRegDefs = make(map[VReg][]programCounter)
+	} else {
+		resetMap(a, i.realRegDefs)
 	}
 }
 

internal/engine/wazevo/backend/regalloc/regalloc_test.go

@@ -446,7 +446,7 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 			for blockID := range a.blockInfos {
 				actual := &a.blockInfos[blockID]
 				exp := tc.exp[blockID]
-				exp.init()
+				initMapInInfo(exp)
 				require.Equal(t, exp, actual, "\n[exp for block[%d]]\n%s\n[actual for block[%d]]\n%s", blockID, exp, blockID, actual)
 			}
 
@@ -560,7 +560,7 @@ func TestAllocator_buildLiveRangesForNonReals(t *testing.T) {
 			require.Equal(t, len(tc.exps), a.nodePool.Allocated())
 			for v, exp := range tc.exps {
 				liveNodes := tc.info.liveNodes
-				tc.info.init()
+				initMapInInfo(tc.info)
 				tc.info.liveNodes = liveNodes
 				t.Run(v.String(), func(t *testing.T) {
 					n := a.vRegIDToNode[v.ID()]
@@ -609,7 +609,7 @@ func TestAllocator_buildLiveRangesForReals(t *testing.T) {
 	} {
 		tc := tc
 		t.Run(tc.name, func(t *testing.T) {
-			tc.info.init()
+			initMapInInfo(tc.info)
 			a := NewAllocator(&RegisterInfo{})
 			a.allocatableSet = tc.allocatableRealRegs
 			a.buildLiveRangesForReals(blockID, tc.info)
@@ -682,3 +682,27 @@ func TestAllocator_recordCopyRelation(t *testing.T) {
 		require.Equal(t, RealRegInvalid, n200.copyToReal)
 	})
 }
+
+func initMapInInfo(info *blockInfo) {
+	if info.liveIns == nil {
+		info.liveIns = make(map[VReg]struct{})
+	}
+	if info.liveOuts == nil {
+		info.liveOuts = make(map[VReg]struct{})
+	}
+	if info.defs == nil {
+		info.defs = make(map[VReg]programCounter)
+	}
+	if info.kills == nil {
+		info.kills = make(map[VReg]programCounter)
+	}
+	if info.lastUses == nil {
+		info.lastUses = make(map[VReg]programCounter)
+	}
+	if info.realRegUses == nil {
+		info.realRegUses = make(map[VReg][]programCounter)
+	}
+	if info.realRegDefs == nil {
+		info.realRegDefs = make(map[VReg][]programCounter)
+	}
+}

internal/engine/wazevo/call_engine.go

@@ -97,7 +97,7 @@ func alignedStackTop(s []byte) uintptr {
 
 // Definition implements api.Function.
 func (c *callEngine) Definition() api.FunctionDefinition {
-	return &c.parent.module.Source.FunctionDefinitionSection[c.indexInModule]
+	return c.parent.module.Source.FunctionDefinition(c.indexInModule)
 }
 
 // Call implements api.Function.

internal/engine/wazevo/e2e_test.go

@@ -34,7 +34,8 @@ func TestSpectestV1(t *testing.T) {
 
 	// TODO: adds incrementally one by one as we support more test cases. And eventually remove this
 	// and migrate to integration_test/spectest/v1/spec_test.go by the time when closing https://github.com/tetratelabs/wazero/issues/1496
-	spectest.RunJson(t, v1.Testcases, "binary.json", context.Background(), config)
+	t.Run("address.json", func(t *testing.T) { spectest.RunJson(t, v1.Testcases, "address.json", context.Background(), config) })
+	t.Run("binary.json", func(t *testing.T) { spectest.RunJson(t, v1.Testcases, "binary.json", context.Background(), config) })
 }
 
 func TestE2E(t *testing.T) {

internal/engine/wazevo/engine.go

@@ -127,6 +127,10 @@ func (e *engine) CompileModule(_ context.Context, module *wasm.Module, _ []exper
 		// Run SSA-level optimization passes.
 		ssaBuilder.RunPasses()
 
+		if false {
+			fmt.Printf("[[[SSA]]]%s", ssaBuilder.Format())
+		}
+
 		// Finalize the layout of SSA blocks which might use the optimization results.
 		ssaBuilder.LayoutBlocks()