Repository: tetratelabs/wazero
Pull request #1763 (2023-10-06 02:11:23)

wazevo(ssa): eliminates no-op shifts by constant zeros
detected by the fuzzer that immediate shift with zero cannot be encoded in arm64.

#1496 


Patch:
diff --git a/internal/engine/wazevo/backend/isa/arm64/instr_encoding.go b/internal/engine/wazevo/backend/isa/arm64/instr_encoding.go
index 7573a64a..41de44ad 100644
--- a/internal/engine/wazevo/backend/isa/arm64/instr_encoding.go
+++ b/internal/engine/wazevo/backend/isa/arm64/instr_encoding.go
@@ -1715,6 +1715,9 @@ func encodeMoveWideImmediate(opc uint32, rd uint32, imm, shift, _64bit uint64) (
 // encodeAluRRImm encodes as "Bitfield" in
 // https://developer.arm.com/documentation/ddi0596/2020-12/Index-by-Encoding/Data-Processing----Immediate?lang=en#log_imm
 func encodeAluRRImm(op aluOp, rd, rn, amount, _64bit uint32) uint32 {
+	if amount == 0 {
+		panic("Shifting by 0 is not allowed; must be optimized out")
+	}
 	var opc uint32
 	var immr, imms uint32
 	switch op {
diff --git a/internal/engine/wazevo/ssa/instructions.go b/internal/engine/wazevo/ssa/instructions.go
index 7b08d00e..0fe82ec5 100644
--- a/internal/engine/wazevo/ssa/instructions.go
+++ b/internal/engine/wazevo/ssa/instructions.go
@@ -1769,11 +1769,12 @@ func (i *Instruction) AsVUshr(x, amount Value, lane VecLane) *Instruction {
 }
 
 // AsSshr initializes this instruction as an integer signed shift right (arithmetic shift right) instruction with OpcodeSshr.
-func (i *Instruction) AsSshr(x, amount Value) {
+func (i *Instruction) AsSshr(x, amount Value) *Instruction {
 	i.opcode = OpcodeSshr
 	i.v = x
 	i.v2 = amount
 	i.typ = x.Type()
+	return i
 }
 
 // AsVSshr initializes this instruction as an integer signed shift right (arithmetic shift right) instruction with OpcodeVSshr on vector.
diff --git a/internal/engine/wazevo/ssa/pass.go b/internal/engine/wazevo/ssa/pass.go
index 07d10f73..146a7416 100644
--- a/internal/engine/wazevo/ssa/pass.go
+++ b/internal/engine/wazevo/ssa/pass.go
@@ -17,6 +17,7 @@ func (b *builder) RunPasses() {
 	passRedundantPhiEliminationOpt(b)
 	// The result of passCalculateImmediateDominators will be used by various passes below.
 	passCalculateImmediateDominators(b)
+	passNopInstElimination(b)
 
 	// TODO: implement either conversion of irreducible CFG into reducible one, or irreducible CFG detection where we panic.
 	// 	WebAssembly program shouldn't result in irreducible CFG, but we should handle it properly in just in case.
@@ -303,3 +304,49 @@ func (b *builder) clearBlkVisited() {
 	}
 	b.blkStack2 = b.blkStack2[:0]
 }
+
+// passNopInstElimination eliminates the instructions which is essentially a no-op.
+func passNopInstElimination(b *builder) {
+	if int(b.nextValueID) >= len(b.valueIDToInstruction) {
+		b.valueIDToInstruction = append(b.valueIDToInstruction, make([]*Instruction, b.nextValueID)...)
+	}
+
+	for blk := b.blockIteratorBegin(); blk != nil; blk = b.blockIteratorNext() {
+		for cur := blk.rootInstr; cur != nil; cur = cur.next {
+			r1, rs := cur.Returns()
+			if r1.Valid() {
+				b.valueIDToInstruction[r1.ID()] = cur
+			}
+			for _, r := range rs {
+				b.valueIDToInstruction[r.ID()] = cur
+			}
+		}
+	}
+
+	for blk := b.blockIteratorBegin(); blk != nil; blk = b.blockIteratorNext() {
+		for cur := blk.rootInstr; cur != nil; cur = cur.next {
+			switch cur.Opcode() {
+			// TODO: add more logics here.
+			case OpcodeIshl, OpcodeSshr, OpcodeUshr:
+				x, amount := cur.Arg2()
+				definingInst := b.valueIDToInstruction[amount.ID()]
+				if definingInst == nil {
+					// If there's no defining instruction, that means the amount is coming from the parameter.
+					continue
+				}
+				if definingInst.Constant() {
+					v := definingInst.ConstantVal()
+
+					if x.Type().Bits() == 64 {
+						v = v % 64
+					} else {
+						v = v % 32
+					}
+					if v == 0 {
+						b.alias(cur.Return(), x)
+					}
+				}
+			}
+		}
+	}
+}


Test patch:
diff --git a/internal/engine/wazevo/backend/backend_test.go b/internal/engine/wazevo/backend/backend_test.go
index 4f06c5b2..3722aa99 100644
--- a/internal/engine/wazevo/backend/backend_test.go
+++ b/internal/engine/wazevo/backend/backend_test.go
@@ -35,7 +35,7 @@ func newMachine() backend.Machine {
 }
 
 func TestE2E(t *testing.T) {
-	const verbose = true
+	const verbose = false
 
 	type testCase struct {
 		name                                   string
diff --git a/internal/engine/wazevo/frontend/frontend_test.go b/internal/engine/wazevo/frontend/frontend_test.go
index 1d7fc9f8..f50d3218 100644
--- a/internal/engine/wazevo/frontend/frontend_test.go
+++ b/internal/engine/wazevo/frontend/frontend_test.go
@@ -1,7 +1,6 @@
 package frontend
 
 import (
-	"fmt"
 	"testing"
 
 	"github.com/tetratelabs/wazero/api"
@@ -1809,13 +1808,11 @@ blk0: (exec_ctx:i64, module_ctx:i64, v2:v128, v3:v128)
 			fc.LowerToSSA()
 
 			actual := fc.formatBuilder()
-			fmt.Println(actual)
 			require.Equal(t, tc.exp, actual)
 
 			b.RunPasses()
 			if expAfterOpt := tc.expAfterOpt; expAfterOpt != "" {
 				actualAfterOpt := fc.formatBuilder()
-				fmt.Println(actualAfterOpt)
 				require.Equal(t, expAfterOpt, actualAfterOpt)
 			}
 
diff --git a/internal/engine/wazevo/ssa/pass_test.go b/internal/engine/wazevo/ssa/pass_test.go
index 7d22debb..2a831a67 100644
--- a/internal/engine/wazevo/ssa/pass_test.go
+++ b/internal/engine/wazevo/ssa/pass_test.go
@@ -1,7 +1,6 @@
 package ssa
 
 import (
-	"fmt"
 	"testing"
 
 	"github.com/tetratelabs/wazero/internal/testing/require"
@@ -11,7 +10,9 @@ func TestBuilder_passes(t *testing.T) {
 	for _, tc := range []struct {
 		name string
 		// pass is the optimization pass to run.
-		pass func(b *builder)
+		pass,
+		// postPass is run after the pass is executed, and can be used to test a pass that depends on another pass.
+		postPass func(b *builder)
 		// setup creates the SSA function in the given *builder.
 		// TODO: when we have the text SSA IR parser, we can eliminate this `setup`,
 		// 	we could directly decode the *builder from the `before` string. I am still
@@ -281,6 +282,58 @@ blk0: ()
 blk1: () <-- (blk0)
 	v4:i32 = Iadd v2, v0
 	Return v4
+`,
+		},
+		{
+			name:     "nop elimination",
+			pass:     passNopInstElimination,
+			postPass: passDeadCodeEliminationOpt,
+			setup: func(b *builder) (verifier func(t *testing.T)) {
+				entry := b.AllocateBasicBlock()
+				b.SetCurrentBlock(entry)
+
+				i32Param := entry.AddParam(b, TypeI32)
+				i64Param := entry.AddParam(b, TypeI64)
+
+				// 32-bit shift.
+				moduleZeroI32 := b.AllocateInstruction().AsIconst32(32 * 245).Insert(b).Return()
+				nopIshl := b.AllocateInstruction().AsIshl(i32Param, moduleZeroI32).Insert(b).Return()
+
+				// 64-bit shift.
+				moduleZeroI64 := b.AllocateInstruction().AsIconst64(64 * 245).Insert(b).Return()
+				nopUshr := b.AllocateInstruction().AsUshr(i64Param, moduleZeroI64).Insert(b).Return()
+
+				// Non zero shift amount should not be eliminated.
+				nonZeroI32 := b.AllocateInstruction().AsIconst32(32*245 + 1).Insert(b).Return()
+				nonZeroIshl := b.AllocateInstruction().AsIshl(i32Param, nonZeroI32).Insert(b).Return()
+
+				nonZeroI64 := b.AllocateInstruction().AsIconst64(64*245 + 1).Insert(b).Return()
+				nonZeroSshr := b.AllocateInstruction().AsSshr(i64Param, nonZeroI64).Insert(b).Return()
+
+				ret := b.AllocateInstruction()
+				ret.AsReturn([]Value{nopIshl, nopUshr, nonZeroIshl, nonZeroSshr})
+				b.InsertInstruction(ret)
+				return nil
+			},
+			before: `
+blk0: (v0:i32, v1:i64)
+	v2:i32 = Iconst_32 0x1ea0
+	v3:i32 = Ishl v0, v2
+	v4:i64 = Iconst_64 0x3d40
+	v5:i64 = Ushr v1, v4
+	v6:i32 = Iconst_32 0x1ea1
+	v7:i32 = Ishl v0, v6
+	v8:i64 = Iconst_64 0x3d41
+	v9:i64 = Sshr v1, v8
+	Return v3, v5, v7, v9
+`,
+			after: `
+blk0: (v0:i32, v1:i64)
+	v6:i32 = Iconst_32 0x1ea1
+	v7:i32 = Ishl v0, v6
+	v8:i64 = Iconst_64 0x3d41
+	v9:i64 = Sshr v1, v8
+	Return v0, v1, v7, v9
 `,
 		},
 	} {
@@ -288,14 +341,14 @@ blk1: () <-- (blk0)
 		t.Run(tc.name, func(t *testing.T) {
 			b := NewBuilder().(*builder)
 			verifier := tc.setup(b)
-			fmt.Println(b.Format())
 			require.Equal(t, tc.before, b.Format())
 			tc.pass(b)
-			fmt.Println("--------")
-			fmt.Println(b.Format())
 			if verifier != nil {
 				verifier(t)
 			}
+			if tc.postPass != nil {
+				tc.postPass(b)
+			}
 			require.Equal(t, tc.after, b.Format())
 		})
 	}