Repository: tetratelabs/wazero
Pull request #2117 (2024-03-05 01:35:28)

wazevo(frontend): stricter BCE
## sqlite speedtest binary size
amd64: 8478767 bytes on main.
amd64: 8337167 bytes with predecessor check.
amd64: 8345407 bytes without predecessor check.

arm64: 8624468 on main.
arm64: 8465156 bytes with predecessor check.
arm64: 8474788 bytes without predecessor check.

### result
amd64: 8345407 - 8337167 = 141600 bytes: ~1.6% of 8.3MB
arm64: 8624468 - 8465156 = 159312 bytes: ~1.8% of 8.6MB

## zig stdlib binary size

amd64: 47128091 bytes on main.
amd64: 45244091 bytes with predecessor check.

arm64: 48182088 bytes on main.
arm64: 46195880 bytes with predecessor check.

### result
amd64: 47128091 - 45244091 = 1884000 bytes: ~4.0% of 47MB
arm64: 48182088 - 46195880 = 1986208 bytes: ~4.1% of 48MB


mathetake: https://github.com/tetratelabs/wazero/pull/2118 will fix the fuzzing failure, which is not relevant to this

Patch:
diff --git a/internal/engine/wazevo/frontend/frontend.go b/internal/engine/wazevo/frontend/frontend.go
index 2ec53406..4e466d5f 100644
--- a/internal/engine/wazevo/frontend/frontend.go
+++ b/internal/engine/wazevo/frontend/frontend.go
@@ -59,8 +59,11 @@ type Compiler struct {
 	execCtxPtrValue, moduleCtxPtrValue ssa.Value
 }
 
+// knownSafeBound represents a known safe bound for a value.
 type knownSafeBound struct {
-	bound        uint64
+	// bound is a constant upper bound for the value.
+	bound uint64
+	// absoluteAddr is the absolute address of the value.
 	absoluteAddr ssa.Value
 }
 
@@ -203,6 +206,7 @@ func (c *Compiler) Init(idx, typIndex wasm.Index, typ *wasm.FunctionType, localT
 	c.wasmFunctionBody = body
 	c.wasmFunctionBodyOffsetInCodeSection = bodyOffsetInCodeSection
 	c.needListener = needListener
+	c.clearSafeBounds()
 }
 
 // Note: this assumes 64-bit platform (I believe we won't have 32-bit backend ;)).
@@ -436,8 +440,7 @@ func (c *Compiler) recordKnownSafeBound(v ssa.ValueID, safeBound uint64, absolut
 	}
 }
 
-// clearSafeBounds clears the known safe bounds. This must be called
-// after the compilation of each block.
+// clearSafeBounds clears the known safe bounds.
 func (c *Compiler) clearSafeBounds() {
 	for _, v := range c.knownSafeBoundsSet {
 		ptr := &c.knownSafeBounds[v]
@@ -446,6 +449,14 @@ func (c *Compiler) clearSafeBounds() {
 	c.knownSafeBoundsSet = c.knownSafeBoundsSet[:0]
 }
 
+// resetAbsoluteAddressInSafeBounds resets the absolute addresses recorded in the known safe bounds.
+func (c *Compiler) resetAbsoluteAddressInSafeBounds() {
+	for _, v := range c.knownSafeBoundsSet {
+		ptr := &c.knownSafeBounds[v]
+		ptr.absoluteAddr = ssa.ValueInvalid
+	}
+}
+
 func (k *knownSafeBound) valid() bool {
 	return k != nil && k.bound > 0
 }
diff --git a/internal/engine/wazevo/frontend/lower.go b/internal/engine/wazevo/frontend/lower.go
index 0674ff90..2e8ce6ae 100644
--- a/internal/engine/wazevo/frontend/lower.go
+++ b/internal/engine/wazevo/frontend/lower.go
@@ -1415,7 +1415,10 @@ func (c *Compiler) lowerCurrentOpcode() {
 		builder.Seal(thenBlk)
 		builder.Seal(elseBlk)
 	case wasm.OpcodeElse:
-		c.clearSafeBounds() // Reset the safe bounds since we are entering the Else block.
+		// Reset the safe bounds since we are entering the Else block.
+		// TODO: we should be able to inherit the safe bounds from the parent block. So, right now, this means that
+		//  else block is a little bit more slow than the then block.
+		c.clearSafeBounds()
 
 		ifctrl := state.ctrlPeekAt(0)
 		if unreachable := state.unreachable; unreachable && state.unreachableDepth > 0 {
@@ -1444,8 +1447,6 @@ func (c *Compiler) lowerCurrentOpcode() {
 		builder.SetCurrentBlock(elseBlk)
 
 	case wasm.OpcodeEnd:
-		c.clearSafeBounds() // Reset the safe bounds since we are exiting the block.
-
 		if state.unreachableDepth > 0 {
 			state.unreachableDepth--
 			break
@@ -1479,7 +1480,14 @@ func (c *Compiler) lowerCurrentOpcode() {
 			c.insertJumpToBlock(ctrl.clonedArgs, followingBlk)
 		}
 
-		builder.Seal(ctrl.followingBlock)
+		builder.Seal(followingBlk)
+
+		if unreachable || followingBlk.Preds() != 1 {
+			// If we can reach this block without being unreachable, and it has only one predecessor,
+			// this means that we get here from the unique block contiguously. Therefore, we can
+			// keep using the same safe bounds information. Otherwise, we need to reset it.
+			c.clearSafeBounds()
+		}
 
 		// Ready to start translating the following block.
 		c.switchTo(ctrl.originalStackLenWithoutParam, followingBlk)
@@ -3641,6 +3649,18 @@ func (c *Compiler) memOpSetup(baseAddr ssa.Value, constOffset, operationSizeInBy
 		// We reuse the calculated absolute address even if the bound is not known to be safe.
 		address = known.absoluteAddr
 		if ceil <= known.bound {
+			if !address.Valid() {
+				// This means that, the bound is known to be safe, but the memory base might have changed.
+				// So, we re-calculate the address.
+				memBase := c.getMemoryBaseValue(false)
+				extBaseAddr := builder.AllocateInstruction().
+					AsUExtend(baseAddr, 32, 64).
+					Insert(builder).
+					Return()
+				address = builder.AllocateInstruction().
+					AsIadd(memBase, extBaseAddr).Insert(builder).Return()
+				known.absoluteAddr = address // Update the absolute address for the subsequent memory access.
+			}
 			return
 		}
 	}
@@ -3761,8 +3781,9 @@ func (c *Compiler) reloadMemoryBaseLen() {
 	_ = c.getMemoryLenValue(true)
 
 	// This function being called means that the memory base might have changed.
-	// Therefore, we need to clear the known safe bounds because we cache the absolute address of the memory access per each base offset.
-	c.clearSafeBounds()
+	// Therefore, we need to clear the absolute addresses recorded in the known safe bounds
+	// because we cache the absolute address of the memory access per each base offset.
+	c.resetAbsoluteAddressInSafeBounds()
 }
 
 func (c *Compiler) setWasmGlobalValue(index wasm.Index, v ssa.Value) {


Test patch:
diff --git a/internal/engine/wazevo/frontend/frontend_test.go b/internal/engine/wazevo/frontend/frontend_test.go
index 9341bd9b..51fbc422 100644
--- a/internal/engine/wazevo/frontend/frontend_test.go
+++ b/internal/engine/wazevo/frontend/frontend_test.go
@@ -3004,6 +3004,29 @@ func TestCompiler_clearSafeBounds(t *testing.T) {
 	require.Equal(t, []knownSafeBound{{}, {}, {}, {}, {}, {}}, c.knownSafeBounds)
 }
 
+func TestCompiler_resetAbsoluteAddressInSafeBounds(t *testing.T) {
+	c := &Compiler{}
+	c.knownSafeBounds = []knownSafeBound{
+		{bound: 1, absoluteAddr: ssa.Value(1)},
+		{},
+		{bound: 2, absoluteAddr: ssa.Value(2)},
+		{},
+		{},
+		{bound: 3, absoluteAddr: ssa.Value(3)},
+	}
+	c.knownSafeBoundsSet = []ssa.ValueID{0, 2, 5}
+	c.resetAbsoluteAddressInSafeBounds()
+	require.Equal(t, 3, len(c.knownSafeBoundsSet))
+	require.Equal(t, []knownSafeBound{
+		{bound: 1, absoluteAddr: ssa.ValueInvalid},
+		{},
+		{bound: 2, absoluteAddr: ssa.ValueInvalid},
+		{},
+		{},
+		{bound: 3, absoluteAddr: ssa.ValueInvalid},
+	}, c.knownSafeBounds)
+}
+
 func TestKnownSafeBound_valid(t *testing.T) {
 	k := &knownSafeBound{bound: 10, absoluteAddr: 12345}
 	require.True(t, k.valid())