wazevo: passes: br_table, loads, select, store, and other spectests (#…

…1658)

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

internal/engine/wazevo/backend/backend_test.go

@@ -2797,57 +2797,73 @@ L1 (SSA Block: blk0):
 	orr w9?, wzr, #0x6
 	subs wzr, w2?, w9?
 	csel w10?, w9?, w2?, hs
-	br_table_sequence x10?, [L2, L3, L4, L5, L6, L7, L2]
-L2 (SSA Block: blk6):
-	movz w0, #0xb, LSL 0
-	ret
-L3 (SSA Block: blk5):
+	br_table_sequence x10?, [L2, L3, L4, L5, L6, L7, L8]
+L2 (SSA Block: blk7):
+	b L9
+L3 (SSA Block: blk8):
+L10 (SSA Block: blk5):
 	orr w0, wzr, #0xc
 	ret
-L4 (SSA Block: blk4):
+L4 (SSA Block: blk9):
+L11 (SSA Block: blk4):
 	movz w0, #0xd, LSL 0
 	ret
-L5 (SSA Block: blk3):
+L5 (SSA Block: blk10):
+L12 (SSA Block: blk3):
 	orr w0, wzr, #0xe
 	ret
-L6 (SSA Block: blk2):
+L6 (SSA Block: blk11):
+L13 (SSA Block: blk2):
 	orr w0, wzr, #0xf
 	ret
-L7 (SSA Block: blk1):
+L7 (SSA Block: blk12):
+L14 (SSA Block: blk1):
 	orr w0, wzr, #0x10
 	ret
+L8 (SSA Block: blk13):
+L9 (SSA Block: blk6):
+	movz w0, #0xb, LSL 0
+	ret
 `,
 			afterFinalizeARM64: `
 L1 (SSA Block: blk0):
 	str x30, [sp, #-0x10]!
 	orr w8, wzr, #0x6
 	subs wzr, w2, w8
 	csel w8, w8, w2, hs
-	adr x27, #16; ldrsw x8, [x27, x8, UXTW 2]; add x27, x27, x8; br x27; [0x1c 0x28 0x34 0x40 0x4c 0x58 0x1c]
-L2 (SSA Block: blk6):
-	movz w0, #0xb, LSL 0
-	ldr x30, [sp], #0x10
-	ret
-L3 (SSA Block: blk5):
+	adr x27, #16; ldrsw x8, [x27, x8, UXTW 2]; add x27, x27, x8; br x27; [0x1c 0x20 0x2c 0x38 0x44 0x50 0x5c]
+L2 (SSA Block: blk7):
+	b #0x40 (L9)
+L3 (SSA Block: blk8):
+L10 (SSA Block: blk5):
 	orr w0, wzr, #0xc
 	ldr x30, [sp], #0x10
 	ret
-L4 (SSA Block: blk4):
+L4 (SSA Block: blk9):
+L11 (SSA Block: blk4):
 	movz w0, #0xd, LSL 0
 	ldr x30, [sp], #0x10
 	ret
-L5 (SSA Block: blk3):
+L5 (SSA Block: blk10):
+L12 (SSA Block: blk3):
 	orr w0, wzr, #0xe
 	ldr x30, [sp], #0x10
 	ret
-L6 (SSA Block: blk2):
+L6 (SSA Block: blk11):
+L13 (SSA Block: blk2):
 	orr w0, wzr, #0xf
 	ldr x30, [sp], #0x10
 	ret
-L7 (SSA Block: blk1):
+L7 (SSA Block: blk12):
+L14 (SSA Block: blk1):
 	orr w0, wzr, #0x10
 	ldr x30, [sp], #0x10
 	ret
+L8 (SSA Block: blk13):
+L9 (SSA Block: blk6):
+	movz w0, #0xb, LSL 0
+	ldr x30, [sp], #0x10
+	ret
 `,
 		},
 	} {

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

@@ -209,6 +209,17 @@ func (m *machine) LowerInstr(instr *ssa.Instruction) {
 		cnt := m.allocateInstr()
 		cnt.asVecMisc(vecOpCvt32To64, rd, rn, vecArrangementNone)
 		m.insert(cnt)
+	case ssa.OpcodeIreduce:
+		rn := m.getOperand_NR(m.compiler.ValueDefinition(instr.UnaryData()), extModeNone)
+		retVal := instr.Return()
+		rd := m.compiler.VRegOf(retVal)
+
+		if retVal.Type() != ssa.TypeI32 {
+			panic("TODO?: Ireduce to non-i32")
+		}
+		mov := m.allocateInstr()
+		mov.asMove32(rd, rn.reg())
+		m.insert(mov)
 	default:
 		panic("TODO: lowering " + instr.Opcode().String())
 	}

internal/engine/wazevo/e2e_test.go

@@ -39,6 +39,7 @@ func TestSpectestV1(t *testing.T) {
 	}{
 		{name: "address"},
 		{name: "br"},
+		{name: "br_table"},
 		{name: "break-drop"},
 		{name: "block"},
 		{name: "binary"},
@@ -47,10 +48,21 @@ func TestSpectestV1(t *testing.T) {
 		{name: "custom"},
 		{name: "const"},
 		{name: "data"},
-		{name: "local_get"},
+		{name: "elem"},
+		{name: "exports"},
+		{name: "globals"},
+		{name: "if"},
+		{name: "names"},
+		{name: "load"},
 		{name: "memory_size"},
 		{name: "memory_grow"},
 		{name: "nop"},
+		{name: "return"},
+		{name: "select"},
+		{name: "store"},
+		{name: "type"},
+		{name: "unreachable"},
+		{name: "unreached-invalid"},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
 			spectest.RunCase(t, v1.Testcases, tc.name, context.Background(), config,

internal/engine/wazevo/engine.go

@@ -108,7 +108,7 @@ func (e *engine) CompileModule(_ context.Context, module *wasm.Module, _ []exper
 		const debug = false
 		if debug {
 			name := exportedFnIndex[fidx]
-			fmt.Printf("%[1]s %d/%d %s %[1]s\n", strings.Repeat("-", 10), i, len(module.CodeSection)-1, name)
+			fmt.Printf("%[1]s %d/%d %s %[1]s\n", strings.Repeat("-", 30), i, len(module.CodeSection)-1, name)
 		}
 
 		_, needGoEntryPreamble := exportedFnIndex[fidx]

internal/engine/wazevo/frontend/frontend_test.go

@@ -1464,31 +1464,52 @@ blk0: (exec_ctx:i64, module_ctx:i64, v2:i32)
 			name: "br_table", m: testcases.BrTable.Module,
 			exp: `
 blk0: (exec_ctx:i64, module_ctx:i64, v2:i32)
-	BrTable v2, [blk6, blk5, blk4, blk3, blk2, blk1, blk6]
+	BrTable v2, [blk7, blk8, blk9, blk10, blk11, blk12, blk13]
 
-blk1: () <-- (blk0)
+blk1: () <-- (blk12)
 	v8:i32 = Iconst_32 0x10
 	Return v8
 
-blk2: () <-- (blk0)
+blk2: () <-- (blk11)
 	v7:i32 = Iconst_32 0xf
 	Return v7
 
-blk3: () <-- (blk0)
+blk3: () <-- (blk10)
 	v6:i32 = Iconst_32 0xe
 	Return v6
 
-blk4: () <-- (blk0)
+blk4: () <-- (blk9)
 	v5:i32 = Iconst_32 0xd
 	Return v5
 
-blk5: () <-- (blk0)
+blk5: () <-- (blk8)
 	v4:i32 = Iconst_32 0xc
 	Return v4
 
-blk6: () <-- (blk0,blk0)
+blk6: () <-- (blk7,blk13)
 	v3:i32 = Iconst_32 0xb
 	Return v3
+
+blk7: () <-- (blk0)
+	Jump blk6
+
+blk8: () <-- (blk0)
+	Jump blk5
+
+blk9: () <-- (blk0)
+	Jump blk4
+
+blk10: () <-- (blk0)
+	Jump blk3
+
+blk11: () <-- (blk0)
+	Jump blk2
+
+blk12: () <-- (blk0)
+	Jump blk1
+
+blk13: () <-- (blk0)
+	Jump blk6
 `,
 		},
 		{

internal/engine/wazevo/frontend/lower.go

@@ -1561,47 +1561,28 @@ func (c *Compiler) lowerBrTable(labels []uint32, index ssa.Value) {
 	}
 
 	targets := make([]ssa.BasicBlock, len(labels))
-	if numArgs == 0 {
-		for i, l := range labels {
-			targetBlk, argNum := state.brTargetArgNumFor(l)
-			if argNum != 0 {
-				// This must be handled in else block below.
-				panic("BUG: br_table with args must not reach here")
-			}
-			targets[i] = targetBlk
 
-			if targetBlk.ReturnBlock() {
-				// TODO: even when the target block has no arguments, we have to insert the unconditional jump to the return trampoline
-				//  if the target is return.
-				panic("TODO")
-			}
-		}
-	} else {
-		// If this needs to pass arguments, we need trampoline blocks since depending on the target block structure,
-		// we might end up inserting moves before jumps, which cannot be done with br_table. Instead, we can do such
-		// per-block moves in the trampoline blocks.
-
-		args := c.loweringState.nPeekDup(numArgs) // Args are always on the top of the stack.
-		currentBlk := builder.CurrentBlock()
-		for i, l := range labels {
-			targetBlk, _ := state.brTargetArgNumFor(l)
-			trampoline := builder.AllocateBasicBlock()
-			builder.SetCurrentBlock(trampoline)
-			c.insertJumpToBlock(args, targetBlk)
-			targets[i] = trampoline
-		}
-		builder.SetCurrentBlock(currentBlk)
+	// We need trampoline blocks since depending on the target block structure, we might end up inserting moves before jumps,
+	// which cannot be done with br_table. Instead, we can do such per-block moves in the trampoline blocks.
+	// At the linking phase (very end of the backend), we can remove the unnecessary jumps, and therefore no runtime overhead.
+	args := c.loweringState.nPeekDup(numArgs) // Args are always on the top of the stack.
+	currentBlk := builder.CurrentBlock()
+	for i, l := range labels {
+		targetBlk, _ := state.brTargetArgNumFor(l)
+		trampoline := builder.AllocateBasicBlock()
+		builder.SetCurrentBlock(trampoline)
+		c.insertJumpToBlock(args, targetBlk)
+		targets[i] = trampoline
 	}
+	builder.SetCurrentBlock(currentBlk)
 
 	// If the target block has no arguments, we can just jump to the target block.
 	brTable := builder.AllocateInstruction()
 	brTable.AsBrTable(index, targets)
 	builder.InsertInstruction(brTable)
 
-	if numArgs > 0 {
-		for _, trampoline := range targets {
-			builder.Seal(trampoline)
-		}
+	for _, trampoline := range targets {
+		builder.Seal(trampoline)
 	}
 }
 

internal/engine/wazevo/ssa/instructions.go

@@ -865,6 +865,7 @@ var instructionSideEffects = [opcodeEnd]sideEffect{
 	OpcodeFcvtFromSint:       sideEffectFalse,
 	OpcodeFcvtFromUint:       sideEffectFalse,
 	OpcodeFpromote:           sideEffectFalse,
+	OpcodeIreduce:            sideEffectFalse,
 }
 
 // HasSideEffects returns true if this instruction has side effects.
@@ -887,6 +888,7 @@ var instructionReturnTypes = [opcodeEnd]returnTypesFn{
 	OpcodeSelect:    returnTypesFnSingle,
 	OpcodeSExtend:   returnTypesFnSingle,
 	OpcodeUExtend:   returnTypesFnSingle,
+	OpcodeIreduce:   returnTypesFnSingle,
 	OpcodeCallIndirect: func(b *builder, instr *Instruction) (t1 Type, ts []Type) {
 		sigID := SignatureID(instr.v)
 		sig, ok := b.signatures[sigID]
@@ -1552,7 +1554,8 @@ func (i *Instruction) Format(b Builder) string {
 	case OpcodeIshl, OpcodeSshr, OpcodeUshr:
 		instSuffix = fmt.Sprintf(" %s, %s", i.v.Format(b), i.v2.Format(b))
 	case OpcodeUndefined:
-	case OpcodeClz, OpcodeCtz, OpcodePopcnt, OpcodeFneg, OpcodeFcvtFromSint, OpcodeFcvtFromUint, OpcodeFpromote:
+	case OpcodeClz, OpcodeCtz, OpcodePopcnt, OpcodeFneg, OpcodeFcvtFromSint, OpcodeFcvtFromUint, OpcodeFpromote,
+		OpcodeIreduce:
 		instSuffix = " " + i.v.Format(b)
 	default:
 		panic(fmt.Sprintf("TODO: format for %s", i.opcode))
@@ -1582,7 +1585,7 @@ func (i *Instruction) addArgumentBranchInst(v Value) {
 	case OpcodeJump, OpcodeBrz, OpcodeBrnz:
 		i.vs = append(i.vs, v)
 	default:
-		panic("BUG: " + i.typ.String())
+		panic("BUG: " + i.opcode.String())
 	}
 }