Repository: tetratelabs/wazero
Pull request #2164 (2024-03-26 10:34:09)

wazevo(frontend): BCE should not propagate the memory address to unsealed blocks
tl;dr: I propose to propagate the memory address only for **sealed** blocks, where we know that no predecessors will be added at a later time.

This originated from an issue raised by @jerbob92 for go-pdfium.

Essentially, it looks like block 4 of function `5011` (local ID = 4960) in go-pdfium references some values that *may* not be visibile on by some of its predecessors.

We are propagating known bounds _while_ we are creating basic blocks; which means that, at a give time, we _might not know_ all of the possible predecessors to a given block. In the case of that block 4, the SSA dump shows:

```
blk4: (v305:i32,v341:i32) <-- (blk3,blk62)
```

which means that one of the predecessors in `blk62`. However, when we enter function `initializeCurrentBlockKnownBounds()` we enter `case 1` because there is only one predecessor (`blk3`) at that time. 

When blocks have more than 1 predecessor, we should enter the `default` case instead: this never propagates `kb.bound` AND `kb.address`, but it always propagates `kb.bound` and `ssa.ValueInvalid` because it cannot assume that the address will be valid on all paths.

This small change, as opposed to other earlier attempts, where I was unconditionally propagating `ssa.ValueInvalid` in all cases (see the Slack threads), does not affect codegen in a lot of cases (e.g. test case `TestCompiler_LowerToSSA/bounds check in if-else` is left unaffected)



mathetake: this is exactly what I was thinking is the right way to fix the issue, so looks good

Patch:
diff --git a/internal/engine/wazevo/entrypoint_amd64.go b/internal/engine/wazevo/entrypoint_amd64.go
index 6cb0c457..18f60af3 100644
--- a/internal/engine/wazevo/entrypoint_amd64.go
+++ b/internal/engine/wazevo/entrypoint_amd64.go
@@ -1,3 +1,5 @@
+//go:build amd64 && !tinygo
+
 package wazevo
 
 import _ "unsafe"
diff --git a/internal/engine/wazevo/entrypoint_arm64.go b/internal/engine/wazevo/entrypoint_arm64.go
index 0b8b5533..e16d64f6 100644
--- a/internal/engine/wazevo/entrypoint_arm64.go
+++ b/internal/engine/wazevo/entrypoint_arm64.go
@@ -1,3 +1,5 @@
+//go:build arm64 && !tinygo
+
 package wazevo
 
 import _ "unsafe"
diff --git a/internal/engine/wazevo/entrypoint_others.go b/internal/engine/wazevo/entrypoint_others.go
index f6d3a6b3..8f9d64b2 100644
--- a/internal/engine/wazevo/entrypoint_others.go
+++ b/internal/engine/wazevo/entrypoint_others.go
@@ -1,4 +1,4 @@
-//go:build !arm64 && !amd64
+//go:build (!arm64 && !amd64) || tinygo
 
 package wazevo
 
diff --git a/internal/engine/wazevo/frontend/frontend.go b/internal/engine/wazevo/frontend/frontend.go
index 8e5a5c3f..873a35a5 100644
--- a/internal/engine/wazevo/frontend/frontend.go
+++ b/internal/engine/wazevo/frontend/frontend.go
@@ -524,7 +524,13 @@ func (c *Compiler) initializeCurrentBlockKnownBounds() {
 	case 1:
 		pred := currentBlk.Pred(0).ID()
 		for _, kb := range c.getKnownSafeBoundsAtTheEndOfBlocks(pred).View() {
-			c.recordKnownSafeBound(kb.id, kb.bound, kb.absoluteAddr)
+			// Unless the block is sealed, we cannot assume the absolute address is valid:
+			// later we might add another predecessor that has no visibility of that value.
+			addr := ssa.ValueInvalid
+			if currentBlk.Sealed() {
+				addr = kb.absoluteAddr
+			}
+			c.recordKnownSafeBound(kb.id, kb.bound, addr)
 		}
 	default:
 		c.pointers = c.pointers[:0]
diff --git a/internal/engine/wazevo/ssa/basic_block.go b/internal/engine/wazevo/ssa/basic_block.go
index 49e4c0f8..10b6b4b6 100644
--- a/internal/engine/wazevo/ssa/basic_block.go
+++ b/internal/engine/wazevo/ssa/basic_block.go
@@ -55,6 +55,9 @@ type BasicBlock interface {
 	// Valid is true if this block is still valid even after optimizations.
 	Valid() bool
 
+	// Sealed is true if this block has been sealed.
+	Sealed() bool
+
 	// BeginPredIterator returns the first predecessor of this block.
 	BeginPredIterator() BasicBlock
 
@@ -212,6 +215,11 @@ func (bb *basicBlock) Valid() bool {
 	return !bb.invalid
 }
 
+// Sealed implements BasicBlock.Sealed.
+func (bb *basicBlock) Sealed() bool {
+	return bb.sealed
+}
+
 // InsertInstruction implements BasicBlock.InsertInstruction.
 func (bb *basicBlock) InsertInstruction(next *Instruction) {
 	current := bb.currentInstr
diff --git a/internal/platform/cpuid_amd64.go b/internal/platform/cpuid_amd64.go
index 2ae784d9..8c9f1a9f 100644
--- a/internal/platform/cpuid_amd64.go
+++ b/internal/platform/cpuid_amd64.go
@@ -1,3 +1,5 @@
+//go:build amd64 && !tinygo
+
 package platform
 
 // CpuFeatures exposes the capabilities for this CPU, queried via the Has, HasExtra methods
diff --git a/internal/platform/cpuid_unsupported.go b/internal/platform/cpuid_unsupported.go
index 953567ce..8ae826d3 100644
--- a/internal/platform/cpuid_unsupported.go
+++ b/internal/platform/cpuid_unsupported.go
@@ -1,4 +1,4 @@
-//go:build !amd64
+//go:build !amd64 || tinygo
 
 package platform
 
diff --git a/internal/platform/mmap_unix.go b/internal/platform/mmap_unix.go
index 72c52150..f29bdc3f 100644
--- a/internal/platform/mmap_unix.go
+++ b/internal/platform/mmap_unix.go
@@ -1,4 +1,4 @@
-//go:build darwin || linux || freebsd
+//go:build (darwin || linux || freebsd) && !tinygo
 
 package platform
 
diff --git a/internal/platform/mmap_unsupported.go b/internal/platform/mmap_unsupported.go
index e1b9f896..4bcbbde8 100644
--- a/internal/platform/mmap_unsupported.go
+++ b/internal/platform/mmap_unsupported.go
@@ -1,4 +1,4 @@
-//go:build !(darwin || linux || freebsd || windows)
+//go:build !(darwin || linux || freebsd || windows) || tinygo
 
 package platform
 
diff --git a/internal/platform/mremap_other.go b/internal/platform/mremap_other.go
index 5be501c1..5cba99fb 100644
--- a/internal/platform/mremap_other.go
+++ b/internal/platform/mremap_other.go
@@ -1,4 +1,4 @@
-//go:build !(darwin || linux || freebsd)
+//go:build !(darwin || linux || freebsd) || tinygo
 
 package platform
 
diff --git a/internal/platform/mremap_unix.go b/internal/platform/mremap_unix.go
index 5017aaf7..8f42d44f 100644
--- a/internal/platform/mremap_unix.go
+++ b/internal/platform/mremap_unix.go
@@ -1,4 +1,4 @@
-//go:build darwin || linux || freebsd
+//go:build (darwin || linux || freebsd) && !tinygo
 
 package platform
 


Test patch:
diff --git a/internal/engine/wazevo/frontend/frontend_test.go b/internal/engine/wazevo/frontend/frontend_test.go
index 39d0be5b..a02543f5 100644
--- a/internal/engine/wazevo/frontend/frontend_test.go
+++ b/internal/engine/wazevo/frontend/frontend_test.go
@@ -2891,6 +2891,105 @@ blk3: () <-- (blk1,blk2)
 	v21:i64 = Iadd v19, v20
 	v22:i32 = Load v21, 0x15
 	Jump blk_ret
+`,
+		},
+		{
+			name: "bounds check in loop with if-else",
+			m: &wasm.Module{
+				TypeSection:     []wasm.FunctionType{{Params: []wasm.ValueType{wasm.ValueTypeI32}}},
+				FunctionSection: []wasm.Index{0},
+				CodeSection: []wasm.Code{{
+					Body: []byte{
+						wasm.OpcodeLocalGet, 0,
+						wasm.OpcodeI32Load, 0x2, 0x10, // staticOffset=0x20
+						wasm.OpcodeDrop,
+
+						wasm.OpcodeLoop, 0x40, // blockSignature:vv.
+
+						/* */ wasm.OpcodeLocalGet, 0,
+						// This address must be always recomputed:
+						// v9 in `Load v9, 0x20` won't propagate from blk0 to blk1,
+						/* */ wasm.OpcodeI32Load, 0x2, 0x20, // staticOffset=0x30
+
+						/* */ wasm.OpcodeIf, 0x40, // blockSignature:vv.
+						/* */ /* */ wasm.OpcodeLoop, 0x40, // blockSignature:vv.
+						/* */ /* */ /* */ wasm.OpcodeLocalGet, 0,
+						/* */ /* */ /* */ wasm.OpcodeI32Load, 0x2, 0x50, // staticOffset=0x50
+						/* */ /* */ /* */ wasm.OpcodeBrIf, 2, // quit to the inner loop
+						/* */ /* */ wasm.OpcodeEnd,
+						/* */ wasm.OpcodeElse,
+						/* */ /* */ wasm.OpcodeLocalGet, 0,
+						/* */ /* */ wasm.OpcodeBrIf, 2, // quit to the outer loop
+						/* */ wasm.OpcodeEnd,
+						/* */ wasm.OpcodeBr, 1, // quit to the loop
+
+						wasm.OpcodeEnd, // end loop
+
+						wasm.OpcodeEnd,
+					},
+				}},
+				MemorySection: &wasm.Memory{Min: 1},
+			},
+			features: api.CoreFeaturesV2,
+			exp: `
+blk0: (exec_ctx:i64, module_ctx:i64, v2:i32)
+	v3:i64 = Iconst_64 0x14
+	v4:i64 = UExtend v2, 32->64
+	v5:i64 = Uload32 module_ctx, 0x10
+	v6:i64 = Iadd v4, v3
+	v7:i32 = Icmp lt_u, v5, v6
+	ExitIfTrue v7, exec_ctx, memory_out_of_bounds
+	v8:i64 = Load module_ctx, 0x8
+	v9:i64 = Iadd v8, v4
+	v10:i32 = Load v9, 0x10
+	Jump blk1
+
+blk1: () <-- (blk0,blk6)
+	v11:i64 = Iconst_64 0x24
+	v12:i64 = UExtend v2, 32->64
+	v13:i64 = Uload32 module_ctx, 0x10
+	v14:i64 = Iadd v12, v11
+	v15:i32 = Icmp lt_u, v13, v14
+	ExitIfTrue v15, exec_ctx, memory_out_of_bounds
+	v16:i64 = Load module_ctx, 0x8
+	v17:i64 = Iadd v16, v12
+	v18:i32 = Load v17, 0x20
+	Brz v18, blk4
+	Jump blk3
+
+blk2: ()
+
+blk3: () <-- (blk1)
+	Jump blk6
+
+blk4: () <-- (blk1)
+	Brnz v2, blk_ret
+	Jump blk9
+
+blk5: () <-- (blk7,blk9)
+	Jump blk_ret
+
+blk6: () <-- (blk3)
+	v19:i64 = Iconst_64 0x54
+	v20:i64 = UExtend v2, 32->64
+	v21:i64 = Uload32 module_ctx, 0x10
+	v22:i64 = Iadd v20, v19
+	v23:i32 = Icmp lt_u, v21, v22
+	ExitIfTrue v23, exec_ctx, memory_out_of_bounds
+	v24:i64 = Load module_ctx, 0x8
+	v25:i64 = Iadd v24, v20
+	v26:i32 = Load v25, 0x50
+	Brnz v26, blk1
+	Jump blk8
+
+blk7: () <-- (blk8)
+	Jump blk5
+
+blk8: () <-- (blk6)
+	Jump blk7
+
+blk9: () <-- (blk4)
+	Jump blk5
 `,
 		},
 	} {
@@ -3136,7 +3235,7 @@ func TestCompiler_finalizeKnownSafeBoundsAtTheEndOoBlock(t *testing.T) {
 }
 
 func TestCompiler_initializeCurrentBlockKnownBounds(t *testing.T) {
-	t.Run("single", func(t *testing.T) {
+	t.Run("single (sealed)", func(t *testing.T) {
 		c := NewFrontendCompiler(&wasm.Module{}, ssa.NewBuilder(), nil, false, false, false)
 		builder := c.ssaBuilder
 		child := builder.AllocateBasicBlock()
@@ -3150,6 +3249,8 @@ func TestCompiler_initializeCurrentBlockKnownBounds(t *testing.T) {
 			c.finalizeKnownSafeBoundsAtTheEndOfBlock(parent.ID())
 		}
 
+		// Seal the child: the absolute addresses will be propagated.
+		builder.Seal(child)
 		builder.SetCurrentBlock(child)
 		c.initializeCurrentBlockKnownBounds()
 		kb := c.getKnownSafeBound(1)
@@ -3165,6 +3266,36 @@ func TestCompiler_initializeCurrentBlockKnownBounds(t *testing.T) {
 		require.Equal(t, uint64(666), kb.bound)
 		require.Equal(t, ssa.Value(54321), kb.absoluteAddr)
 	})
+	t.Run("single (unsealed)", func(t *testing.T) {
+		c := NewFrontendCompiler(&wasm.Module{}, ssa.NewBuilder(), nil, false, false, false)
+		builder := c.ssaBuilder
+		child := builder.AllocateBasicBlock()
+		{
+			parent := builder.AllocateBasicBlock()
+			builder.SetCurrentBlock(parent)
+			c.recordKnownSafeBound(1, 99, 9999)
+			c.recordKnownSafeBound(2, 150, 9999)
+			c.recordKnownSafeBound(5, 666, 54321)
+			builder.AllocateInstruction().AsJump(ssa.ValuesNil, child).Insert(builder)
+			c.finalizeKnownSafeBoundsAtTheEndOfBlock(parent.ID())
+		}
+
+		// Do not seal the child: the absolute addresses will be recomputed.
+		builder.SetCurrentBlock(child)
+		c.initializeCurrentBlockKnownBounds()
+		kb := c.getKnownSafeBound(1)
+		require.True(t, kb.valid())
+		require.Equal(t, uint64(99), kb.bound)
+		require.NotEqual(t, ssa.Value(9999), kb.absoluteAddr)
+		kb = c.getKnownSafeBound(2)
+		require.True(t, kb.valid())
+		require.Equal(t, uint64(150), kb.bound)
+		require.NotEqual(t, ssa.Value(9999), kb.absoluteAddr)
+		kb = c.getKnownSafeBound(5)
+		require.True(t, kb.valid())
+		require.Equal(t, uint64(666), kb.bound)
+		require.NotEqual(t, ssa.Value(54321), kb.absoluteAddr)
+	})
 	t.Run("multiple predecessors", func(t *testing.T) {
 		c := NewFrontendCompiler(&wasm.Module{}, ssa.NewBuilder(), nil, false, false, false)
 		builder := c.ssaBuilder