wazevo(amd64): lower Icmp(Band(x,y), 0) to one TEST (tetratelabs#2073)

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

internal/engine/wazevo/backend/backend_test.go

@@ -7,6 +7,8 @@ import (
 	"runtime"
 	"testing"
 
+	"github.com/tetratelabs/wazero/api"
+	"github.com/tetratelabs/wazero/experimental"
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend"
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend/isa/amd64"
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend/isa/arm64"
@@ -45,7 +47,7 @@ func TestE2E(t *testing.T) {
 		m                                      *wasm.Module
 		targetIndex                            uint32
 		afterLoweringARM64, afterFinalizeARM64 string
-		// TODO: amd64.
+		afterLoweringAMD64, afterFinalizeAMD64 string
 	}
 
 	for _, tc := range []testCase{
@@ -2362,6 +2364,46 @@ L2:
 	ldr x30, [sp], #0x10
 	add sp, sp, #0x10
 	ret
+`,
+		},
+		{
+			name: "icmp_and_zero",
+			m:    testcases.IcmpAndZero.Module,
+			afterFinalizeAMD64: `
+L1 (SSA Block: blk0):
+	pushq %rbp
+	movq %rsp, %rbp
+	testl %edi, %ecx
+	jnz L2
+L3 (SSA Block: blk1):
+	movl $1, %eax
+	movq %rbp, %rsp
+	popq %rbp
+	ret
+L2 (SSA Block: blk2):
+	xor %rax, %rax
+	movq %rbp, %rsp
+	popq %rbp
+	ret
+`,
+			afterFinalizeARM64: `
+L1 (SSA Block: blk0):
+	stp x30, xzr, [sp, #-0x10]!
+	str xzr, [sp, #-0x10]!
+	ands wzr, w2, w3
+	b.ne #0x18, (L2)
+L3 (SSA Block: blk1):
+	orr w8, wzr, #0x1
+	mov x0, x8
+	add sp, sp, #0x10
+	ldr x30, [sp], #0x10
+	ret
+L2 (SSA Block: blk2):
+	mov x8, xzr
+	mov x0, x8
+	add sp, sp, #0x10
+	ldr x30, [sp], #0x10
+	ret
 `,
 		},
 	} {
@@ -2371,11 +2413,18 @@ L2:
 			case "arm64":
 				exp = tc.afterFinalizeARM64
 			case "amd64":
-				t.Skip()
+				if tc.afterFinalizeAMD64 != "" {
+					exp = tc.afterFinalizeAMD64
+				} else {
+					t.Skip()
+				}
 			default:
 				t.Fail()
 			}
 
+			err := tc.m.Validate(api.CoreFeaturesV2 | experimental.CoreFeaturesThreads)
+			require.NoError(t, err)
+
 			ssab := ssa.NewBuilder()
 			offset := wazevoapi.NewModuleContextOffsetData(tc.m, false)
 			fc := frontend.NewFrontendCompiler(tc.m, ssab, &offset, false, false, false)
@@ -2418,6 +2467,10 @@ L2:
 				if tc.afterLoweringARM64 != "" {
 					require.Equal(t, tc.afterLoweringARM64, be.Format())
 				}
+			case "amd64":
+				if tc.afterLoweringAMD64 != "" {
+					require.Equal(t, tc.afterLoweringAMD64, be.Format())
+				}
 			default:
 				t.Fail()
 			}

internal/engine/wazevo/backend/isa/amd64/machine.go

@@ -286,7 +286,9 @@ func (m *machine) LowerConditionalBranch(b *ssa.Instruction) {
 
 		// First, perform the comparison and set the flag.
 		xd, yd := m.c.ValueDefinition(x), m.c.ValueDefinition(y)
-		m.lowerIcmpToFlag(xd, yd, x.Type() == ssa.TypeI64)
+		if !m.tryLowerBandToFlag(xd, yd) {
+			m.lowerIcmpToFlag(xd, yd, x.Type() == ssa.TypeI64)
+		}
 
 		// Then perform the conditional branch.
 		m.insert(m.allocateInstr().asJmpIf(cc, newOperandLabel(target)))
@@ -1296,14 +1298,46 @@ func (m *machine) lowerExitIfTrueWithCode(execCtx regalloc.VReg, cond ssa.Value,
 	execCtxTmp := m.copyToTmp(execCtx)
 
 	x, y, c := cvalInstr.IcmpData()
-	m.lowerIcmpToFlag(m.c.ValueDefinition(x), m.c.ValueDefinition(y), x.Type() == ssa.TypeI64)
+	xx, yy := m.c.ValueDefinition(x), m.c.ValueDefinition(y)
+	if !m.tryLowerBandToFlag(xx, yy) {
+		m.lowerIcmpToFlag(xx, yy, x.Type() == ssa.TypeI64)
+	}
 
 	jmpIf := m.allocateInstr()
 	m.insert(jmpIf)
 	l := m.lowerExitWithCode(execCtxTmp, code)
 	jmpIf.asJmpIf(condFromSSAIntCmpCond(c).invert(), newOperandLabel(l))
 }
 
+func (m *machine) tryLowerBandToFlag(x, y *backend.SSAValueDefinition) (ok bool) {
+	var target *backend.SSAValueDefinition
+	if x.IsFromInstr() && x.Instr.Constant() && x.Instr.ConstantVal() == 0 {
+		if m.c.MatchInstr(y, ssa.OpcodeBand) {
+			target = y
+		}
+	}
+
+	if y.IsFromInstr() && y.Instr.Constant() && y.Instr.ConstantVal() == 0 {
+		if m.c.MatchInstr(x, ssa.OpcodeBand) {
+			target = x
+		}
+	}
+
+	if target == nil {
+		return false
+	}
+
+	bandInstr := target.Instr
+	bandX, bandY := bandInstr.Arg2()
+
+	xx := m.getOperand_Reg(m.c.ValueDefinition(bandX))
+	yy := m.getOperand_Mem_Imm32_Reg(m.c.ValueDefinition(bandY))
+	test := m.allocateInstr().asCmpRmiR(false, yy, xx.reg(), bandX.Type() == ssa.TypeI64)
+	m.insert(test)
+	bandInstr.MarkLowered()
+	return true
+}
+
 func (m *machine) allocateExitInstructions(execCtx, exitCodeReg regalloc.VReg) (saveRsp, saveRbp, setExitCode *instruction) {
 	saveRsp = m.allocateInstr().asMovRM(
 		rspVReg,

internal/engine/wazevo/testcases/testcases.go

@@ -2379,6 +2379,30 @@ var (
 			}}},
 		},
 	}
+
+	IcmpAndZero = TestCase{
+		Name: "icmp_and_zero",
+		Module: &wasm.Module{
+			TypeSection:     []wasm.FunctionType{{Params: []wasm.ValueType{i32, i32}, Results: []wasm.ValueType{i32}}},
+			ExportSection:   []wasm.Export{{Name: ExportedFunctionName, Type: wasm.ExternTypeFunc, Index: 0}},
+			FunctionSection: []wasm.Index{0},
+			CodeSection: []wasm.Code{{Body: []byte{
+				wasm.OpcodeLocalGet, 0,
+				wasm.OpcodeLocalGet, 1,
+				wasm.OpcodeI32And,
+				wasm.OpcodeI32Eqz,
+				wasm.OpcodeIf, blockSignature_vv,
+				wasm.OpcodeI32Const, 1,
+				wasm.OpcodeReturn,
+				wasm.OpcodeElse,
+				wasm.OpcodeI32Const, 0,
+				wasm.OpcodeReturn,
+				wasm.OpcodeEnd,
+				wasm.OpcodeUnreachable,
+				wasm.OpcodeEnd,
+			}}},
+		},
+	}
 )
 
 // VecShuffleWithLane returns a VecShuffle test with a custom 16-bytes immediate (lane indexes).