Repository: tetratelabs/wazero
Pull request #2109 (2024-03-04 00:27:00)

wazevo(ssa): recursive constant phi eliminations

mathetake: This makes the entire perf really better like 2x faster in some cases. Especially, the sqlite speedtest executable size is like 30% smaller... amazing...
mathetake: fuzzing time

Patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr.go b/internal/engine/wazevo/backend/isa/amd64/instr.go
index bed379ea..600ae2b1 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr.go
@@ -258,6 +258,8 @@ func (i *instruction) String() string {
 		return fmt.Sprintf("xmmcmov%s %s, %s", cond(i.u1), i.op1.format(true), i.op2.format(true))
 	case blendvpd:
 		return fmt.Sprintf("blendvpd %s, %s, %%xmm0", i.op1.format(false), i.op2.format(false))
+	case mfence:
+		return "mfence"
 	default:
 		panic(fmt.Sprintf("BUG: %d", int(i.kind)))
 	}
@@ -754,9 +756,17 @@ const (
 	// blendvpd is https://www.felixcloutier.com/x86/blendvpd.
 	blendvpd
 
+	// mfence is https://www.felixcloutier.com/x86/mfence
+	mfence
+
 	instrMax
 )
 
+func (i *instruction) asMFence() *instruction {
+	i.kind = mfence
+	return i
+}
+
 func (i *instruction) asIdivRemSequence(execCtx, divisor, tmpGp regalloc.VReg, isDiv, signed, _64 bool) *instruction {
 	i.kind = idivRemSequence
 	i.op1 = newOperandReg(execCtx)
@@ -978,6 +988,8 @@ func (k instructionKind) String() string {
 		return "xmmCMov"
 	case idivRemSequence:
 		return "idivRemSequence"
+	case mfence:
+		return "mfence"
 	default:
 		panic("BUG")
 	}
@@ -2218,6 +2230,7 @@ var defKinds = [instrMax]defKind{
 	xmmCMov:                defKindOp2,
 	idivRemSequence:        defKindDivRem,
 	blendvpd:               defKindNone,
+	mfence:                 defKindNone,
 }
 
 // String implements fmt.Stringer.
@@ -2293,6 +2306,7 @@ var useKinds = [instrMax]useKind{
 	xmmCMov:                useKindOp1,
 	idivRemSequence:        useKindDivRem,
 	blendvpd:               useKindBlendvpd,
+	mfence:                 useKindNone,
 }
 
 func (u useKind) String() string {
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go b/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
index 95de5e8e..2f99e8b5 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
@@ -1259,6 +1259,12 @@ func (i *instruction) encode(c backend.Compiler) (needsLabelResolution bool) {
 		}
 		i.encode(c)
 
+	case mfence:
+		// https://www.felixcloutier.com/x86/mfence
+		c.EmitByte(0x0f)
+		c.EmitByte(0xae)
+		c.EmitByte(0xf0)
+
 	default:
 		panic(fmt.Sprintf("TODO: %v", i.kind))
 	}
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine.go b/internal/engine/wazevo/backend/isa/amd64/machine.go
index 26985995..29575ee0 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine.go
@@ -933,6 +933,32 @@ func (m *machine) LowerInstr(instr *ssa.Instruction) {
 		}
 		m.insert(m.allocateInstr().asMovzxRmR(extModeLQ, rn, rd))
 
+	case ssa.OpcodeAtomicLoad:
+		ptr := instr.Arg()
+		size := instr.AtomicTargetSize()
+		dst := m.c.VRegOf(instr.Return())
+
+		// At this point, the ptr is ensured to be aligned, so using a normal load is atomic.
+		// https://github.com/golang/go/blob/adead1a93f472affa97c494ef19f2f492ee6f34a/src/runtime/internal/atomic/atomic_amd64.go#L30
+		mem := newOperandMem(m.lowerToAddressMode(ptr, 0))
+		load := m.allocateInstr()
+		switch size {
+		case 8:
+			load.asMov64MR(mem, dst)
+		case 4:
+			load.asMovzxRmR(extModeLQ, mem, dst)
+		case 2:
+			load.asMovzxRmR(extModeWQ, mem, dst)
+		case 1:
+			load.asMovzxRmR(extModeBQ, mem, dst)
+		default:
+			panic("BUG")
+		}
+		m.insert(load)
+
+	case ssa.OpcodeFence:
+		m.insert(m.allocateInstr().asMFence())
+
 	default:
 		panic("TODO: lowering " + op.String())
 	}
diff --git a/internal/engine/wazevo/ssa/pass.go b/internal/engine/wazevo/ssa/pass.go
index 69afa7cb..a2e986cd 100644
--- a/internal/engine/wazevo/ssa/pass.go
+++ b/internal/engine/wazevo/ssa/pass.go
@@ -112,88 +112,105 @@ func passDeadBlockEliminationOpt(b *builder) {
 func passRedundantPhiEliminationOpt(b *builder) {
 	redundantParameterIndexes := b.ints[:0] // reuse the slice from previous iterations.
 
-	_ = b.blockIteratorBegin() // skip entry block!
-	// Below, we intentionally use the named iteration variable name, as this comes with inevitable nested for loops!
-	for blk := b.blockIteratorNext(); blk != nil; blk = b.blockIteratorNext() {
-		paramNum := len(blk.params)
+	// TODO: this might be costly for large programs, but at least, as far as I did the experiment, it's almost the
+	//  same as the single iteration version in terms of the overall compilation time. That *might be* mostly thanks to the fact
+	//  that removing many PHIs results in the reduction of the total instructions, not because of this indefinite iteration is
+	//  relatively small. For example, sqlite speedtest binary results in the large number of redundant PHIs,
+	//  the maximum number of iteration was 22, which seems to be acceptable but not that small either since the
+	//  complexity here is O(BlockNum * Iterations) at the worst case where BlockNum might be the order of thousands.
+	for {
+		changed := false
+		_ = b.blockIteratorBegin() // skip entry block!
+		// Below, we intentionally use the named iteration variable name, as this comes with inevitable nested for loops!
+		for blk := b.blockIteratorNext(); blk != nil; blk = b.blockIteratorNext() {
+			paramNum := len(blk.params)
+
+			for paramIndex := 0; paramIndex < paramNum; paramIndex++ {
+				phiValue := blk.params[paramIndex].value
+				redundant := true
+
+				nonSelfReferencingValue := ValueInvalid
+				for predIndex := range blk.preds {
+					br := blk.preds[predIndex].branch
+					// Resolve the alias in the arguments so that we could use the previous iteration's result.
+					b.resolveArgumentAlias(br)
+					pred := br.vs.View()[paramIndex]
+					if pred == phiValue {
+						// This is self-referencing: PHI from the same PHI.
+						continue
+					}
 
-		for paramIndex := 0; paramIndex < paramNum; paramIndex++ {
-			phiValue := blk.params[paramIndex].value
-			redundant := true
+					if !nonSelfReferencingValue.Valid() {
+						nonSelfReferencingValue = pred
+						continue
+					}
 
-			nonSelfReferencingValue := ValueInvalid
-			for predIndex := range blk.preds {
-				pred := blk.preds[predIndex].branch.vs.View()[paramIndex]
-				if pred == phiValue {
-					// This is self-referencing: PHI from the same PHI.
-					continue
+					if nonSelfReferencingValue != pred {
+						redundant = false
+						break
+					}
 				}
 
 				if !nonSelfReferencingValue.Valid() {
-					nonSelfReferencingValue = pred
-					continue
+					// This shouldn't happen, and must be a bug in builder.go.
+					panic("BUG: params added but only self-referencing")
 				}
 
-				if nonSelfReferencingValue != pred {
-					redundant = false
-					break
+				if redundant {
+					b.redundantParameterIndexToValue[paramIndex] = nonSelfReferencingValue
+					redundantParameterIndexes = append(redundantParameterIndexes, paramIndex)
 				}
 			}
 
-			if !nonSelfReferencingValue.Valid() {
-				// This shouldn't happen, and must be a bug in builder.go.
-				panic("BUG: params added but only self-referencing")
+			if len(b.redundantParameterIndexToValue) == 0 {
+				continue
 			}
+			changed = true
 
-			if redundant {
-				b.redundantParameterIndexToValue[paramIndex] = nonSelfReferencingValue
-				redundantParameterIndexes = append(redundantParameterIndexes, paramIndex)
+			// Remove the redundant PHIs from the argument list of branching instructions.
+			for predIndex := range blk.preds {
+				var cur int
+				predBlk := blk.preds[predIndex]
+				branchInst := predBlk.branch
+				view := branchInst.vs.View()
+				for argIndex, value := range view {
+					if _, ok := b.redundantParameterIndexToValue[argIndex]; !ok {
+						view[cur] = value
+						cur++
+					}
+				}
+				branchInst.vs.Cut(cur)
 			}
-		}
 
-		if len(b.redundantParameterIndexToValue) == 0 {
-			continue
-		}
+			// Still need to have the definition of the value of the PHI (previously as the parameter).
+			for _, redundantParamIndex := range redundantParameterIndexes {
+				phiValue := blk.params[redundantParamIndex].value
+				onlyValue := b.redundantParameterIndexToValue[redundantParamIndex]
+				// Create an alias in this block from the only phi argument to the phi value.
+				b.alias(phiValue, onlyValue)
+			}
 
-		// Remove the redundant PHIs from the argument list of branching instructions.
-		for predIndex := range blk.preds {
+			// Finally, Remove the param from the blk.
 			var cur int
-			predBlk := blk.preds[predIndex]
-			branchInst := predBlk.branch
-			view := branchInst.vs.View()
-			for argIndex, value := range view {
-				if _, ok := b.redundantParameterIndexToValue[argIndex]; !ok {
-					view[cur] = value
+			for paramIndex := 0; paramIndex < paramNum; paramIndex++ {
+				param := blk.params[paramIndex]
+				if _, ok := b.redundantParameterIndexToValue[paramIndex]; !ok {
+					blk.params[cur] = param
 					cur++
 				}
 			}
-			branchInst.vs.Cut(cur)
-		}
-
-		// Still need to have the definition of the value of the PHI (previously as the parameter).
-		for _, redundantParamIndex := range redundantParameterIndexes {
-			phiValue := blk.params[redundantParamIndex].value
-			onlyValue := b.redundantParameterIndexToValue[redundantParamIndex]
-			// Create an alias in this block from the only phi argument to the phi value.
-			b.alias(phiValue, onlyValue)
-		}
+			blk.params = blk.params[:cur]
 
-		// Finally, Remove the param from the blk.
-		var cur int
-		for paramIndex := 0; paramIndex < paramNum; paramIndex++ {
-			param := blk.params[paramIndex]
-			if _, ok := b.redundantParameterIndexToValue[paramIndex]; !ok {
-				blk.params[cur] = param
-				cur++
+			// Clears the map for the next iteration.
+			for _, paramIndex := range redundantParameterIndexes {
+				delete(b.redundantParameterIndexToValue, paramIndex)
 			}
+			redundantParameterIndexes = redundantParameterIndexes[:0]
 		}
-		blk.params = blk.params[:cur]
 
-		// Clears the map for the next iteration.
-		for _, paramIndex := range redundantParameterIndexes {
-			delete(b.redundantParameterIndexToValue, paramIndex)
+		if !changed {
+			break
 		}
-		redundantParameterIndexes = redundantParameterIndexes[:0]
 	}
 
 	// Reuse the slice for the future passes.
diff --git a/internal/wasm/func_validation.go b/internal/wasm/func_validation.go
index 592ac869..673d5763 100644
--- a/internal/wasm/func_validation.go
+++ b/internal/wasm/func_validation.go
@@ -1090,6 +1090,8 @@ func (m *Module) validateFunctionWithMaxStackValues(
 				for _, r := range results {
 					valueTypeStack.push(r)
 				}
+			} else {
+				return fmt.Errorf("unknown misc opcode %#x", miscOpcode)
 			}
 		} else if op == OpcodeVecPrefix {
 			pc++
@@ -1393,7 +1395,7 @@ func (m *Module) validateFunctionWithMaxStackValues(
 				}
 				valueTypeStack.push(ValueTypeV128)
 			default:
-				return fmt.Errorf("TODO: SIMD instruction %s will be implemented in #506", vectorInstructionName[vecOpcode])
+				return fmt.Errorf("unknown SIMD instruction %s", vectorInstructionName[vecOpcode])
 			}
 		} else if op == OpcodeBlock {
 			br.Reset(body[pc+1:])


Test patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr_encoding_test.go b/internal/engine/wazevo/backend/isa/amd64/instr_encoding_test.go
index b3f9ccf7..be85fa7e 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr_encoding_test.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr_encoding_test.go
@@ -4046,6 +4046,11 @@ func TestInstruction_format_encode(t *testing.T) {
 			want:       "66440f3815f9",
 			wantFormat: "blendvpd %xmm1, %xmm15, %xmm0",
 		},
+		{
+			setup:      func(i *instruction) { i.asMFence() },
+			want:       "0faef0",
+			wantFormat: "mfence",
+		},
 	} {
 		tc := tc
 		t.Run(tc.wantFormat, func(t *testing.T) {
diff --git a/internal/engine/wazevo/e2e_test.go b/internal/engine/wazevo/e2e_test.go
index c0605a41..9807ebbe 100644
--- a/internal/engine/wazevo/e2e_test.go
+++ b/internal/engine/wazevo/e2e_test.go
@@ -600,10 +600,9 @@ func TestE2E(t *testing.T) {
 			},
 		},
 		{
-			name:      "memory_wait32",
-			m:         testcases.MemoryWait32.Module,
-			features:  api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
-			skipAMD64: true,
+			name:     "memory_wait32",
+			m:        testcases.MemoryWait32.Module,
+			features: api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
 			calls: []callCase{
 				{params: []uint64{0x0, 0xbeef, 0xffffffff}, expResults: []uint64{1}}, // exp not equal, returns 1
 				{params: []uint64{0x1, 0xbeef, 0xffffffff}, expErr: "unaligned atomic"},
@@ -615,10 +614,9 @@ func TestE2E(t *testing.T) {
 			},
 		},
 		{
-			name:      "memory_wait64",
-			m:         testcases.MemoryWait64.Module,
-			features:  api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
-			skipAMD64: true,
+			name:     "memory_wait64",
+			m:        testcases.MemoryWait64.Module,
+			features: api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
 			calls: []callCase{
 				{params: []uint64{0x0, 0xbeef, 0xffffffff}, expResults: []uint64{1}}, // exp not equal, returns 1
 				{params: []uint64{0x1, 0xbeef, 0xffffffff}, expErr: "unaligned atomic"},
@@ -634,10 +632,9 @@ func TestE2E(t *testing.T) {
 			},
 		},
 		{
-			name:      "memory_notify",
-			m:         testcases.MemoryNotify.Module,
-			features:  api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
-			skipAMD64: true,
+			name:     "memory_notify",
+			m:        testcases.MemoryNotify.Module,
+			features: api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
 			calls: []callCase{
 				{params: []uint64{0x0, 0x1}, expResults: []uint64{0}}, // no waiters, returns 0
 				{params: []uint64{0x1, 0x1}, expErr: "unaligned atomic"},
@@ -845,10 +842,9 @@ func TestE2E(t *testing.T) {
 			},
 		},
 		{
-			name:      "atomic_fence",
-			m:         testcases.AtomicFence.Module,
-			features:  api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
-			skipAMD64: true,
+			name:     "atomic_fence",
+			m:        testcases.AtomicFence.Module,
+			features: api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
 			calls: []callCase{
 				{params: []uint64{}, expResults: []uint64{}},
 			},
diff --git a/internal/engine/wazevo/ssa/pass_test.go b/internal/engine/wazevo/ssa/pass_test.go
index 261f3225..83d33fbe 100644
--- a/internal/engine/wazevo/ssa/pass_test.go
+++ b/internal/engine/wazevo/ssa/pass_test.go
@@ -189,7 +189,7 @@ blk0: ()
 
 blk1: (v0:i32) <-- (blk0,blk1)
 	v3:i32 = Iconst_32 0xff
-	Brz v2, blk1, v3
+	Brz v1, blk1, v3
 	Jump blk2
 
 blk2: () <-- (blk1)
diff --git a/internal/integration_test/fuzzcases/fuzzcases_test.go b/internal/integration_test/fuzzcases/fuzzcases_test.go
index 3d9cc928..feaa3816 100644
--- a/internal/integration_test/fuzzcases/fuzzcases_test.go
+++ b/internal/integration_test/fuzzcases/fuzzcases_test.go
@@ -1043,3 +1043,16 @@ func Test2097(t *testing.T) {
 	require.Error(t, err)
 	require.Contains(t, err.Error(), "unexpected end of function")
 }
+
+func Test2112(t *testing.T) {
+	bin, err := hex.DecodeString("0061736d0100000001050160017e00020100030201000404017000000a06010400fc300b")
+	require.NoError(t, err)
+	if !platform.CompilerSupported() {
+		return
+	}
+
+	r := wazero.NewRuntimeWithConfig(context.Background(), opt.NewRuntimeConfigOptimizingCompiler())
+	_, err = r.Instantiate(ctx, bin)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "invalid function[0]: unknown misc opcode 0x30")
+}