Repository: tetratelabs/wazero
Pull request #2072 (2024-02-19 05:39:51)

wazevo(arm64): lower Icmp(Band(x,y), 0) to single ANDS instruction
With this patch, arm64 backend can lower `Icmp(Band(x,y), 0)` pattern to 
one ANDS instruction.

This patter appears, for example, the alignment check in threads support and 
the old compiler was handling it better than the optimizing compiler:
https://github.com/tetratelabs/wazero/blob/ccf60cb7a9796fdc52ff771c3bc7e7c72e5a0131/internal/engine/compiler/impl_threads_arm64.go#L40

This patch applies the same optimization to a general lowering logic around branching.
For Zig binary, this amounts approximately up to 1% of the total conditional branches.
In other words, one instruction is optimized out for 1% of all conditional branches.

After this, I will port this to amd64 as well.


Special thanks to @anuraaga 



Patch:
diff --git a/internal/engine/wazevo/backend/isa/arm64/instr.go b/internal/engine/wazevo/backend/isa/arm64/instr.go
index 3aa1279d..10d28223 100644
--- a/internal/engine/wazevo/backend/isa/arm64/instr.go
+++ b/internal/engine/wazevo/backend/isa/arm64/instr.go
@@ -1813,6 +1813,8 @@ func (a aluOp) String() string {
 		return "orn"
 	case aluOpAnd:
 		return "and"
+	case aluOpAnds:
+		return "ands"
 	case aluOpBic:
 		return "bic"
 	case aluOpEor:
@@ -1856,6 +1858,8 @@ const (
 	aluOpOrn
 	// 32/64-bit Bitwise AND.
 	aluOpAnd
+	// 32/64-bit Bitwise ANDS.
+	aluOpAnds
 	// 32/64-bit Bitwise AND NOT.
 	aluOpBic
 	// 32/64-bit Bitwise XOR (Exclusive OR).
diff --git a/internal/engine/wazevo/backend/isa/arm64/instr_encoding.go b/internal/engine/wazevo/backend/isa/arm64/instr_encoding.go
index e654b52c..471a8399 100644
--- a/internal/engine/wazevo/backend/isa/arm64/instr_encoding.go
+++ b/internal/engine/wazevo/backend/isa/arm64/instr_encoding.go
@@ -1376,6 +1376,8 @@ func encodeAluBitmaskImmediate(op aluOp, rd, rn uint32, imm uint64, _64bit bool)
 		_31to23 = 0b01_100100
 	case aluOpEor:
 		_31to23 = 0b10_100100
+	case aluOpAnds:
+		_31to23 = 0b11_100100
 	default:
 		panic("BUG")
 	}
@@ -1539,7 +1541,7 @@ func encodeAluRRRShift(op aluOp, rd, rn, rm, amount uint32, shiftOp shiftOp, _64
 		_31to24 = 0b01001011
 	case aluOpSubS:
 		_31to24 = 0b01101011
-	case aluOpAnd, aluOpOrr, aluOpEor:
+	case aluOpAnd, aluOpOrr, aluOpEor, aluOpAnds:
 		// "Logical (shifted register)".
 		switch op {
 		case aluOpAnd:
@@ -1548,6 +1550,8 @@ func encodeAluRRRShift(op aluOp, rd, rn, rm, amount uint32, shiftOp shiftOp, _64
 			opc = 0b01
 		case aluOpEor:
 			opc = 0b10
+		case aluOpAnds:
+			opc = 0b11
 		}
 		_31to24 = 0b000_01010
 	default:
@@ -1649,7 +1653,7 @@ func encodeAluRRR(op aluOp, rd, rn, rm uint32, _64bit, isRnSp bool) uint32 {
 		}
 		// "Shifted register" with shift = 0
 		_31to21 = 0b01101011_000
-	case aluOpAnd, aluOpOrr, aluOpOrn, aluOpEor:
+	case aluOpAnd, aluOpOrr, aluOpOrn, aluOpEor, aluOpAnds:
 		// "Logical (shifted register)".
 		var opc, n uint32
 		switch op {
@@ -1662,6 +1666,8 @@ func encodeAluRRR(op aluOp, rd, rn, rm uint32, _64bit, isRnSp bool) uint32 {
 			n = 1
 		case aluOpEor:
 			opc = 0b10
+		case aluOpAnds:
+			opc = 0b11
 		}
 		_31to21 = 0b000_01010_000 | opc<<8 | n
 	case aluOpLsl, aluOpAsr, aluOpLsr, aluOpRotR:
diff --git a/internal/engine/wazevo/backend/isa/arm64/lower_instr.go b/internal/engine/wazevo/backend/isa/arm64/lower_instr.go
index 68bd5412..2bb234e8 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_instr.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_instr.go
@@ -90,7 +90,9 @@ func (m *machine) LowerConditionalBranch(b *ssa.Instruction) {
 			cc = cc.invert()
 		}
 
-		m.lowerIcmpToFlag(x, y, signed)
+		if !m.tryLowerBandToFlag(x, y) {
+			m.lowerIcmpToFlag(x, y, signed)
+		}
 		cbr := m.allocateInstr()
 		cbr.asCondBr(cc.asCond(), target, false /* ignored */)
 		m.insert(cbr)
@@ -121,6 +123,31 @@ func (m *machine) LowerConditionalBranch(b *ssa.Instruction) {
 	}
 }
 
+func (m *machine) tryLowerBandToFlag(x, y ssa.Value) (ok bool) {
+	xx := m.compiler.ValueDefinition(x)
+	yy := m.compiler.ValueDefinition(y)
+	if xx.IsFromInstr() && xx.Instr.Constant() && xx.Instr.ConstantVal() == 0 {
+		if m.compiler.MatchInstr(yy, ssa.OpcodeBand) {
+			bandInstr := yy.Instr
+			m.lowerBitwiseAluOp(bandInstr, aluOpAnds, true)
+			ok = true
+			bandInstr.MarkLowered()
+			return
+		}
+	}
+
+	if yy.IsFromInstr() && yy.Instr.Constant() && yy.Instr.ConstantVal() == 0 {
+		if m.compiler.MatchInstr(xx, ssa.OpcodeBand) {
+			bandInstr := xx.Instr
+			m.lowerBitwiseAluOp(bandInstr, aluOpAnds, true)
+			ok = true
+			bandInstr.MarkLowered()
+			return
+		}
+	}
+	return
+}
+
 // LowerInstr implements backend.Machine.
 func (m *machine) LowerInstr(instr *ssa.Instruction) {
 	if l := instr.SourceOffset(); l.Valid() {
@@ -179,11 +206,11 @@ func (m *machine) LowerInstr(instr *ssa.Instruction) {
 	case ssa.OpcodeVMinPseudo:
 		m.lowerVMinMaxPseudo(instr, false)
 	case ssa.OpcodeBand:
-		m.lowerBitwiseAluOp(instr, aluOpAnd)
+		m.lowerBitwiseAluOp(instr, aluOpAnd, false)
 	case ssa.OpcodeBor:
-		m.lowerBitwiseAluOp(instr, aluOpOrr)
+		m.lowerBitwiseAluOp(instr, aluOpOrr, false)
 	case ssa.OpcodeBxor:
-		m.lowerBitwiseAluOp(instr, aluOpEor)
+		m.lowerBitwiseAluOp(instr, aluOpEor, false)
 	case ssa.OpcodeIshl:
 		m.lowerShifts(instr, extModeNone, aluOpLsl)
 	case ssa.OpcodeSshr:
@@ -1645,12 +1672,18 @@ func (m *machine) lowerShifts(si *ssa.Instruction, ext extMode, aluOp aluOp) {
 	m.insert(alu)
 }
 
-func (m *machine) lowerBitwiseAluOp(si *ssa.Instruction, op aluOp) {
+func (m *machine) lowerBitwiseAluOp(si *ssa.Instruction, op aluOp, ignoreResult bool) {
 	x, y := si.Arg2()
 
 	xDef, yDef := m.compiler.ValueDefinition(x), m.compiler.ValueDefinition(y)
 	rn := m.getOperand_NR(xDef, extModeNone)
-	rd := operandNR(m.compiler.VRegOf(si.Return()))
+
+	var rd operand
+	if ignoreResult {
+		rd = operandNR(xzrVReg)
+	} else {
+		rd = operandNR(m.compiler.VRegOf(si.Return()))
+	}
 
 	_64 := x.Type().Bits() == 64
 	alu := m.allocateInstr()
@@ -1934,7 +1967,10 @@ func (m *machine) lowerExitIfTrueWithCode(execCtxVReg regalloc.VReg, cond ssa.Va
 	cvalInstr := condDef.Instr
 	x, y, c := cvalInstr.IcmpData()
 	signed := c.Signed()
-	m.lowerIcmpToFlag(x, y, signed)
+
+	if !m.tryLowerBandToFlag(x, y) {
+		m.lowerIcmpToFlag(x, y, signed)
+	}
 
 	// We need to copy the execution context to a temp register, because if it's spilled,
 	// it might end up being reloaded inside the exiting branch.


Test patch:
diff --git a/internal/engine/wazevo/backend/backend_test.go b/internal/engine/wazevo/backend/backend_test.go
index 6c45fe92..e35b65a9 100644
--- a/internal/engine/wazevo/backend/backend_test.go
+++ b/internal/engine/wazevo/backend/backend_test.go
@@ -2155,6 +2155,213 @@ L1 (SSA Block: blk0):
 	ldr q29, [sp], #0x10
 	ldr x30, [sp], #0x10
 	ret
+`,
+		},
+		{
+			name: "AtomicRmwAdd",
+			m:    testcases.AtomicRmwAdd.Module,
+			afterFinalizeARM64: `
+L1 (SSA Block: blk0):
+	orr x27, xzr, #0x10
+	sub sp, sp, x27
+	stp x30, x27, [sp, #-0x10]!
+	str x19, [sp, #-0x10]!
+	orr x27, xzr, #0x10
+	str x27, [sp, #-0x10]!
+	ldr x8, [sp, #0x30]
+	mov x9, xzr
+	uxtw x9, w9
+	add x10, x1, #0x10
+	ldar x10, x10
+	add x11, x9, #0x1
+	subs xzr, x10, x11
+	mov x10, x0
+	b.hs #0x34, (L13)
+	movz x11, #0x4, lsl 0
+	str w11, [x10]
+	mov x11, sp
+	str x11, [x10, #0x38]
+	adr x11, #0x0
+	str x11, [x10, #0x30]
+	exit_sequence x10
+L13:
+	ldr x10, [x1, #0x8]
+	add x9, x10, x9
+	ldaddalb w2, w9, x9
+	orr w11, wzr, #0x8
+	uxtw x11, w11
+	add x12, x1, #0x10
+	ldar x12, x12
+	add x13, x11, #0x2
+	subs xzr, x12, x13
+	mov x12, x0
+	b.hs #0x34, (L12)
+	movz x13, #0x4, lsl 0
+	str w13, [x12]
+	mov x13, sp
+	str x13, [x12, #0x38]
+	adr x13, #0x0
+	str x13, [x12, #0x30]
+	exit_sequence x12
+L12:
+	add x11, x10, x11
+	ands xzr, x11, #0x1
+	mov x12, x0
+	b.eq #0x34, (L11)
+	movz x13, #0x17, lsl 0
+	str w13, [x12]
+	mov x13, sp
+	str x13, [x12, #0x38]
+	adr x13, #0x0
+	str x13, [x12, #0x30]
+	exit_sequence x12
+L11:
+	ldaddalh w3, w11, x11
+	orr w12, wzr, #0x10
+	uxtw x12, w12
+	add x13, x1, #0x10
+	ldar x13, x13
+	add x14, x12, #0x4
+	subs xzr, x13, x14
+	mov x13, x0
+	b.hs #0x34, (L10)
+	movz x14, #0x4, lsl 0
+	str w14, [x13]
+	mov x14, sp
+	str x14, [x13, #0x38]
+	adr x14, #0x0
+	str x14, [x13, #0x30]
+	exit_sequence x13
+L10:
+	add x12, x10, x12
+	ands xzr, x12, #0x3
+	mov x13, x0
+	b.eq #0x34, (L9)
+	movz x14, #0x17, lsl 0
+	str w14, [x13]
+	mov x14, sp
+	str x14, [x13, #0x38]
+	adr x14, #0x0
+	str x14, [x13, #0x30]
+	exit_sequence x13
+L9:
+	ldaddal w4, w12, x12
+	orr w13, wzr, #0x18
+	uxtw x13, w13
+	add x14, x1, #0x10
+	ldar x14, x14
+	add x15, x13, #0x1
+	subs xzr, x14, x15
+	mov x14, x0
+	b.hs #0x34, (L8)
+	movz x15, #0x4, lsl 0
+	str w15, [x14]
+	mov x15, sp
+	str x15, [x14, #0x38]
+	adr x15, #0x0
+	str x15, [x14, #0x30]
+	exit_sequence x14
+L8:
+	add x13, x10, x13
+	ldaddalb w5, w13, x13
+	orr w14, wzr, #0x20
+	uxtw x14, w14
+	add x15, x1, #0x10
+	ldar x15, x15
+	add x16, x14, #0x2
+	subs xzr, x15, x16
+	mov x15, x0
+	b.hs #0x34, (L7)
+	movz x16, #0x4, lsl 0
+	str w16, [x15]
+	mov x16, sp
+	str x16, [x15, #0x38]
+	adr x16, #0x0
+	str x16, [x15, #0x30]
+	exit_sequence x15
+L7:
+	add x14, x10, x14
+	ands xzr, x14, #0x1
+	mov x15, x0
+	b.eq #0x34, (L6)
+	movz x16, #0x17, lsl 0
+	str w16, [x15]
+	mov x16, sp
+	str x16, [x15, #0x38]
+	adr x16, #0x0
+	str x16, [x15, #0x30]
+	exit_sequence x15
+L6:
+	ldaddalh w6, w14, x14
+	movz w15, #0x28, lsl 0
+	uxtw x15, w15
+	add x16, x1, #0x10
+	ldar x16, x16
+	add x17, x15, #0x4
+	subs xzr, x16, x17
+	mov x16, x0
+	b.hs #0x34, (L5)
+	movz x17, #0x4, lsl 0
+	str w17, [x16]
+	mov x17, sp
+	str x17, [x16, #0x38]
+	adr x17, #0x0
+	str x17, [x16, #0x30]
+	exit_sequence x16
+L5:
+	add x15, x10, x15
+	ands xzr, x15, #0x3
+	mov x16, x0
+	b.eq #0x34, (L4)
+	movz x17, #0x17, lsl 0
+	str w17, [x16]
+	mov x17, sp
+	str x17, [x16, #0x38]
+	adr x17, #0x0
+	str x17, [x16, #0x30]
+	exit_sequence x16
+L4:
+	ldaddal w7, w15, x15
+	orr w16, wzr, #0x30
+	uxtw x16, w16
+	add x17, x1, #0x10
+	ldar x17, x17
+	add x19, x16, #0x8
+	subs xzr, x17, x19
+	mov x17, x0
+	b.hs #0x34, (L3)
+	movz x19, #0x4, lsl 0
+	str w19, [x17]
+	mov x19, sp
+	str x19, [x17, #0x38]
+	adr x19, #0x0
+	str x19, [x17, #0x30]
+	exit_sequence x17
+L3:
+	add x10, x10, x16
+	ands xzr, x10, #0x7
+	b.eq #0x34, (L2)
+	movz x16, #0x17, lsl 0
+	str w16, [x0]
+	mov x16, sp
+	str x16, [x0, #0x38]
+	adr x16, #0x0
+	str x16, [x0, #0x30]
+	exit_sequence x0
+L2:
+	ldaddal x8, x8, x10
+	mov x6, x8
+	mov x5, x15
+	mov x4, x14
+	mov x3, x13
+	mov x2, x12
+	mov x1, x11
+	mov x0, x9
+	add sp, sp, #0x10
+	ldr x19, [sp], #0x10
+	ldr x30, [sp], #0x10
+	add sp, sp, #0x10
+	ret
 `,
 		},
 	} {
diff --git a/internal/engine/wazevo/backend/isa/arm64/instr_encoding_test.go b/internal/engine/wazevo/backend/isa/arm64/instr_encoding_test.go
index 75c3c295..e60a469c 100644
--- a/internal/engine/wazevo/backend/isa/arm64/instr_encoding_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/instr_encoding_test.go
@@ -1035,6 +1035,18 @@ func TestInstruction_encode(t *testing.T) {
 		{want: "2030428a", setup: func(i *instruction) {
 			i.asALU(aluOpAnd, operandNR(x0VReg), operandNR(x1VReg), operandSR(x2VReg, 12, shiftOpLSR), true)
 		}},
+		{want: "2000026a", setup: func(i *instruction) {
+			i.asALU(aluOpAnds, operandNR(x0VReg), operandNR(x1VReg), operandNR(x2VReg), false)
+		}},
+		{want: "200002ea", setup: func(i *instruction) {
+			i.asALU(aluOpAnds, operandNR(x0VReg), operandNR(x1VReg), operandNR(x2VReg), true)
+		}},
+		{want: "201002ea", setup: func(i *instruction) {
+			i.asALU(aluOpAnds, operandNR(x0VReg), operandNR(x1VReg), operandSR(x2VReg, 4, shiftOpLSL), true)
+		}},
+		{want: "203042ea", setup: func(i *instruction) {
+			i.asALU(aluOpAnds, operandNR(x0VReg), operandNR(x1VReg), operandSR(x2VReg, 12, shiftOpLSR), true)
+		}},
 		{want: "2000022a", setup: func(i *instruction) {
 			i.asALU(aluOpOrr, operandNR(x0VReg), operandNR(x1VReg), operandNR(x2VReg), false)
 		}},