Repository: tetratelabs/wazero
Pull request #1673 (2023-08-30 03:20:26)

wazevo: introduces the deterministic compilation verifier
this introduces the new debugging option to ensure that the compilation pipeline is deterministic. This is disabled by default but can be enabled when debugging.

#1496 


Patch:
diff --git a/internal/engine/wazevo/backend/compiler.go b/internal/engine/wazevo/backend/compiler.go
index c8533654..d766fcc3 100644
--- a/internal/engine/wazevo/backend/compiler.go
+++ b/internal/engine/wazevo/backend/compiler.go
@@ -1,6 +1,8 @@
 package backend
 
 import (
+	"context"
+	"encoding/hex"
 	"fmt"
 
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend/regalloc"
@@ -39,7 +41,7 @@ type Compiler interface {
 	//
 	// The returned byte slices are the machine code and the relocation information for the machine code.
 	// The caller is responsible for copying them immediately since the compiler may reuse the buffer.
-	Compile() (_ []byte, _ []RelocationInfo, goPreambleSize int, _ error)
+	Compile(ctx context.Context) (_ []byte, _ []RelocationInfo, goPreambleSize int, _ error)
 
 	// Lower lowers the given ssa.Instruction to the machine-specific instructions.
 	Lower()
@@ -142,20 +144,32 @@ type compiler struct {
 }
 
 // Compile implements Compiler.Compile.
-func (c *compiler) Compile() ([]byte, []RelocationInfo, int, error) {
+func (c *compiler) Compile(ctx context.Context) ([]byte, []RelocationInfo, int, error) {
 	c.Lower()
 	if wazevoapi.PrintSSAToBackendIRLowering {
 		fmt.Printf("[[[after lowering]]]%s\n", c.Format())
 	}
+	if wazevoapi.DeterministicCompilationVerifierEnabled {
+		wazevoapi.VerifyOrSetDeterministicCompilationContextValue(ctx, "After lowering to ISA specific IR", c.Format())
+	}
 	c.RegAlloc()
 	if wazevoapi.PrintRegisterAllocated {
 		fmt.Printf("[[[after regalloc]]]%s\n", c.Format())
 	}
+	if wazevoapi.DeterministicCompilationVerifierEnabled {
+		wazevoapi.VerifyOrSetDeterministicCompilationContextValue(ctx, "After Register Allocation", c.Format())
+	}
 	c.Finalize()
 	if wazevoapi.PrintFinalizedMachineCode {
 		fmt.Printf("[[[after finalize]]]%s\n", c.Format())
 	}
+	if wazevoapi.DeterministicCompilationVerifierEnabled {
+		wazevoapi.VerifyOrSetDeterministicCompilationContextValue(ctx, "After Finalization", c.Format())
+	}
 	goPreambleSize := c.Encode()
+	if wazevoapi.DeterministicCompilationVerifierEnabled {
+		wazevoapi.VerifyOrSetDeterministicCompilationContextValue(ctx, "Encoded Machine code", hex.EncodeToString(c.buf))
+	}
 	return c.buf, c.relocations, goPreambleSize, nil
 }
 
diff --git a/internal/engine/wazevo/engine.go b/internal/engine/wazevo/engine.go
index 8a052c07..8e5a81f6 100644
--- a/internal/engine/wazevo/engine.go
+++ b/internal/engine/wazevo/engine.go
@@ -67,13 +67,33 @@ var _ wasm.Engine = (*engine)(nil)
 // NewEngine returns the implementation of wasm.Engine.
 func NewEngine(_ context.Context, _ api.CoreFeatures, _ filecache.Cache) wasm.Engine {
 	e := &engine{compiledModules: make(map[wasm.ModuleID]*compiledModule), refToBinaryOffset: make(map[ssa.FuncRef]int)}
-
 	e.compileBuiltinFunctions()
 	return e
 }
 
 // CompileModule implements wasm.Engine.
-func (e *engine) CompileModule(_ context.Context, module *wasm.Module, _ []experimental.FunctionListener, _ bool) error {
+func (e *engine) CompileModule(ctx context.Context, module *wasm.Module, listeners []experimental.FunctionListener, ensureTermination bool) error {
+	if wazevoapi.DeterministicCompilationVerifierEnabled {
+		ctx = wazevoapi.NewDeterministicCompilationVerifierContext(ctx, len(module.CodeSection))
+	}
+	cm, err := e.compileModule(ctx, module, listeners, ensureTermination)
+	if err != nil {
+		return err
+	}
+	e.addCompiledModule(module, cm)
+
+	if wazevoapi.DeterministicCompilationVerifierEnabled {
+		for i := 0; i < wazevoapi.DeterministicCompilationVerifyingIter; i++ {
+			_, err = e.compileModule(ctx, module, listeners, ensureTermination)
+			if err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
+func (e *engine) compileModule(ctx context.Context, module *wasm.Module, listeners []experimental.FunctionListener, ensureTermination bool) (*compiledModule, error) {
 	e.rels = e.rels[:0]
 	cm := &compiledModule{offsets: wazevoapi.NewModuleContextOffsetData(module)}
 
@@ -83,8 +103,12 @@ func (e *engine) CompileModule(_ context.Context, module *wasm.Module, _ []exper
 
 	importedFns, localFns := int(module.ImportFunctionCount), len(module.FunctionSection)
 	if localFns == 0 {
-		e.addCompiledModule(module, cm)
-		return nil
+		return cm, nil
+	}
+
+	if wazevoapi.DeterministicCompilationVerifierEnabled {
+		// The compilation must be deterministic regardless of the order of functions being compiled.
+		wazevoapi.DeterministicCompilationVerifierRandomizeIndexes(ctx)
 	}
 
 	// TODO: reuse the map to avoid allocation.
@@ -106,84 +130,28 @@ func (e *engine) CompileModule(_ context.Context, module *wasm.Module, _ []exper
 	cm.functionOffsets = make([]compiledFunctionOffset, localFns)
 	bodies := make([][]byte, localFns)
 	for i := range module.CodeSection {
+		if wazevoapi.DeterministicCompilationVerifierEnabled {
+			i = wazevoapi.DeterministicCompilationVerifierGetRandomizedLocalFunctionIndex(ctx, i)
+		}
+
 		fidx := wasm.Index(i + importedFns)
-		fref := frontend.FunctionIndexToFuncRef(fidx)
 
 		_, needGoEntryPreamble := exportedFnIndex[fidx]
 		if sf := module.StartSection; sf != nil && *sf == fidx {
 			needGoEntryPreamble = true
 		}
 
+		body, rels, goPreambleSize, err := e.compileLocalWasmFunction(ctx, module, wasm.Index(i), fidx, needGoEntryPreamble, fe, ssaBuilder, be, listeners, ensureTermination)
+		if err != nil {
+			return nil, fmt.Errorf("compile function %d/%d: %v", i, len(module.CodeSection)-1, err)
+		}
+
 		// Align 16-bytes boundary.
 		totalSize = (totalSize + 15) &^ 15
 		compiledFuncOffset := &cm.functionOffsets[i]
 		compiledFuncOffset.offset = totalSize
 
-		typ := &module.TypeSection[module.FunctionSection[i]]
-
-		codeSeg := &module.CodeSection[i]
-
-		const checkDeterministicSSACompilation = false
-		if checkDeterministicSSACompilation {
-			existingResults := map[string]struct{}{}
-			for i := 0; i < 100; i++ {
-				fe.Init(wasm.Index(i), typ, codeSeg.LocalTypes, codeSeg.Body)
-				be.Init(needGoEntryPreamble)
-
-				// Lower Wasm to SSA.
-				err := fe.LowerToSSA()
-				if err != nil {
-					return fmt.Errorf("wasm->ssa: %v", err)
-				}
-				ssaBuilder.RunPasses()
-				ssaBuilder.LayoutBlocks()
-
-				existingResults[ssaBuilder.Format()] = struct{}{}
-				if len(existingResults) > 1 {
-					for f := range existingResults {
-						fmt.Println(f)
-						fmt.Println("----------------")
-					}
-					panic("BUG: SSA format is not deterministic")
-				}
-			}
-		}
-
-		// Initializes both frontend and backend compilers.
-		fe.Init(wasm.Index(i), typ, codeSeg.LocalTypes, codeSeg.Body)
-		be.Init(needGoEntryPreamble)
-
-		// Lower Wasm to SSA.
-		err := fe.LowerToSSA()
-		if err != nil {
-			return fmt.Errorf("wasm->ssa: %v", err)
-		}
-
-		if wazevoapi.PrintSSA {
-			fmt.Printf("[[[SSA for %d/%d %s]]]%s\n", i, len(module.CodeSection)-1, exportedFnIndex[fidx], ssaBuilder.Format())
-		}
-
-		// Run SSA-level optimization passes.
-		ssaBuilder.RunPasses()
-
-		if wazevoapi.PrintOptimizedSSA {
-			fmt.Printf("[[[optimized SSA for %d/%d %s]]]%s\n", i, len(module.CodeSection)-1, exportedFnIndex[fidx], ssaBuilder.Format())
-		}
-
-		// Finalize the layout of SSA blocks which might use the optimization results.
-		ssaBuilder.LayoutBlocks()
-
-		if wazevoapi.PrintBlockLaidOutSSA {
-			fmt.Printf("[[[laidout SSA for %d/%d %s]]]%s\n", i, len(module.CodeSection)-1, exportedFnIndex[fidx], ssaBuilder.Format())
-		}
-
-		// Now our ssaBuilder contains the necessary information to further lower them to
-		// machine code.
-		body, rels, goPreambleSize, err := be.Compile()
-		if err != nil {
-			return fmt.Errorf("ssa->machine code: %v", err)
-		}
-
+		fref := frontend.FunctionIndexToFuncRef(fidx)
 		e.refToBinaryOffset[fref] = totalSize +
 			// During the relocation, call target needs to be the beginning of function after Go entry preamble.
 			goPreambleSize
@@ -198,10 +166,7 @@ func (e *engine) CompileModule(_ context.Context, module *wasm.Module, _ []exper
 			e.rels = append(e.rels, r)
 		}
 
-		// TODO: optimize as zero copy.
-		copied := make([]byte, len(body))
-		copy(copied, body)
-		bodies[i] = copied
+		bodies[i] = body
 		totalSize += len(body)
 		if wazevoapi.PrintMachineCodeHexPerFunction {
 			fmt.Printf("[[[machine code SSA for %d/%d %s]]]\n%s\n",
@@ -227,7 +192,7 @@ func (e *engine) CompileModule(_ context.Context, module *wasm.Module, _ []exper
 	if runtime.GOARCH == "arm64" {
 		// On arm64, we cannot give all of rwx at the same time, so we change it to exec.
 		if err = platform.MprotectRX(executable); err != nil {
-			return err
+			return nil, err
 		}
 	}
 	e.compiledModules[module.ID] = cm
@@ -235,11 +200,86 @@ func (e *engine) CompileModule(_ context.Context, module *wasm.Module, _ []exper
 
 	// TODO: finalizer.
 
-	e.addCompiledModule(module, cm)
-	return nil
+	return cm, nil
 }
 
-func (e *engine) compileHostModule(module *wasm.Module) error {
+func (e *engine) compileLocalWasmFunction(
+	ctx context.Context,
+	module *wasm.Module,
+	localFunctionIndex,
+	functionIndex wasm.Index,
+	needGoEntryPreamble bool,
+	fe *frontend.Compiler,
+	ssaBuilder ssa.Builder,
+	be backend.Compiler,
+	_ []experimental.FunctionListener, _ bool,
+) (body []byte, rels []backend.RelocationInfo, goPreambleSize int, err error) {
+	typ := &module.TypeSection[module.FunctionSection[localFunctionIndex]]
+	codeSeg := &module.CodeSection[localFunctionIndex]
+
+	// Initializes both frontend and backend compilers.
+	fe.Init(localFunctionIndex, typ, codeSeg.LocalTypes, codeSeg.Body)
+	be.Init(needGoEntryPreamble)
+
+	// Lower Wasm to SSA.
+	err = fe.LowerToSSA()
+	if err != nil {
+		return nil, nil, 0, fmt.Errorf("wasm->ssa: %v", err)
+	}
+
+	if wazevoapi.PrintSSA {
+		def := module.FunctionDefinition(functionIndex)
+		fmt.Printf("[[[SSA for %d/%d %s]]]%s\n", localFunctionIndex, len(module.CodeSection)-1, def.Debugname, ssaBuilder.Format())
+	}
+
+	if wazevoapi.DeterministicCompilationVerifierEnabled {
+		def := module.FunctionDefinition(functionIndex)
+		name := def.DebugName()
+		if len(def.ExportNames()) > 0 {
+			name = def.ExportNames()[0]
+		}
+		ctx = wazevoapi.DeterministicCompilationVerifierSetCurrentFunctionName(ctx, fmt.Sprintf("[%d/%d]:\"%s\"", localFunctionIndex, len(module.CodeSection)-1, name))
+		wazevoapi.VerifyOrSetDeterministicCompilationContextValue(ctx, "SSA", ssaBuilder.Format())
+	}
+
+	// Run SSA-level optimization passes.
+	ssaBuilder.RunPasses()
+
+	if wazevoapi.PrintOptimizedSSA {
+		def := module.FunctionDefinition(functionIndex)
+		fmt.Printf("[[[Optimized SSA for %d/%d %s]]]%s\n", localFunctionIndex, len(module.CodeSection)-1, def.Debugname, ssaBuilder.Format())
+	}
+
+	if wazevoapi.DeterministicCompilationVerifierEnabled {
+		wazevoapi.VerifyOrSetDeterministicCompilationContextValue(ctx, "Optimized SSA", ssaBuilder.Format())
+	}
+
+	// Finalize the layout of SSA blocks which might use the optimization results.
+	ssaBuilder.LayoutBlocks()
+
+	if wazevoapi.PrintBlockLaidOutSSA {
+		def := module.FunctionDefinition(functionIndex)
+		fmt.Printf("[[[Laidout SSA for %d/%d %s]]]%s\n", localFunctionIndex, len(module.CodeSection)-1, def.Debugname, ssaBuilder.Format())
+	}
+
+	if wazevoapi.DeterministicCompilationVerifierEnabled {
+		wazevoapi.VerifyOrSetDeterministicCompilationContextValue(ctx, "Block laid out SSA", ssaBuilder.Format())
+	}
+
+	// Now our ssaBuilder contains the necessary information to further lower them to
+	// machine code.
+	original, rels, goPreambleSize, err := be.Compile(ctx)
+	if err != nil {
+		return nil, nil, 0, fmt.Errorf("ssa->machine code: %v", err)
+	}
+
+	// TODO: optimize as zero copy.
+	copied := make([]byte, len(original))
+	copy(copied, original)
+	return copied, rels, goPreambleSize, nil
+}
+
+func (e *engine) compileHostModule(module *wasm.Module) (*compiledModule, error) {
 	machine := newMachine()
 	be := backend.NewCompiler(machine, ssa.NewBuilder())
 
@@ -256,7 +296,7 @@ func (e *engine) compileHostModule(module *wasm.Module) error {
 		typIndex := module.FunctionSection[i]
 		typ := &module.TypeSection[typIndex]
 		if typ.ParamNumInUint64 >= goFunctionCallStackSize || typ.ResultNumInUint64 >= goFunctionCallStackSize {
-			return fmt.Errorf("too many params or results for a host function (maximum %d): %v",
+			return nil, fmt.Errorf("too many params or results for a host function (maximum %d): %v",
 				goFunctionCallStackSize, typ)
 		}
 
@@ -264,7 +304,7 @@ func (e *engine) compileHostModule(module *wasm.Module) error {
 		// However, 1 << 16 should be large enough for a real use case.
 		const hostFunctionNumMaximum = 1 << 16
 		if i >= hostFunctionNumMaximum {
-			return fmt.Errorf("too many host functions (maximum %d)", hostFunctionNumMaximum)
+			return nil, fmt.Errorf("too many host functions (maximum %d)", hostFunctionNumMaximum)
 		}
 
 		sig.ID = ssa.SignatureID(typIndex) // This is important since we reuse the `machine` which caches the ABI based on the SignatureID.
@@ -322,12 +362,10 @@ func (e *engine) compileHostModule(module *wasm.Module) error {
 	if runtime.GOARCH == "arm64" {
 		// On arm64, we cannot give all of rwx at the same time, so we change it to exec.
 		if err = platform.MprotectRX(executable); err != nil {
-			return err
+			return nil, err
 		}
 	}
-
-	e.compiledModules[module.ID] = cm
-	return nil
+	return cm, nil
 }
 
 // Close implements wasm.Engine.
diff --git a/internal/engine/wazevo/wazevoapi/debug_consts.go b/internal/engine/wazevo/wazevoapi/debug_consts.go
index fa807974..12c6e14c 100644
--- a/internal/engine/wazevo/wazevoapi/debug_consts.go
+++ b/internal/engine/wazevo/wazevoapi/debug_consts.go
@@ -1,5 +1,13 @@
 package wazevoapi
 
+import (
+	"context"
+	"fmt"
+	"math/rand"
+	"os"
+	"time"
+)
+
 // These consts are used various places in the wazevo implementations.
 // Instead of defining them in each file, we define them here so that we can quickly iterate on
 // debugging without spending "where do we have debug logging?" time.
@@ -38,3 +46,91 @@ const (
 	RegAllocValidationEnabled = true
 	SSAValidationEnabled      = true
 )
+
+// ----- Deterministic compilation verifier -----
+
+const (
+	// DeterministicCompilationVerifierEnabled enables the deterministic compilation verifier. This is disabled by default
+	// since the operation is expensive. But when in doubt, enable this to make sure the compilation is deterministic.
+	DeterministicCompilationVerifierEnabled = false
+	DeterministicCompilationVerifyingIter   = 20
+)
+
+type (
+	verifierContext struct {
+		randomizedIndexes []int
+		r                 *rand.Rand
+		values            map[string]string
+	}
+	existingValuesContextKey struct{}
+	currentFunctionNameKey   struct{}
+)
+
+// NewDeterministicCompilationVerifierContext creates a new context with the deterministic compilation verifier used per wasm.Module.
+func NewDeterministicCompilationVerifierContext(ctx context.Context, localFunctions int) context.Context {
+	randomizedIndexes := make([]int, localFunctions)
+	for i := range randomizedIndexes {
+		randomizedIndexes[i] = i
+	}
+	r := rand.New(rand.NewSource(time.Now().UnixNano()))
+	return context.WithValue(ctx, existingValuesContextKey{}, &verifierContext{
+		r: r, randomizedIndexes: randomizedIndexes, values: map[string]string{},
+	})
+}
+
+// DeterministicCompilationVerifierRandomizeIndexes randomizes the indexes for the deterministic compilation verifier.
+// To get the randomized index, use DeterministicCompilationVerifierGetRandomizedLocalFunctionIndex.
+func DeterministicCompilationVerifierRandomizeIndexes(ctx context.Context) {
+	verifierCtx := ctx.Value(existingValuesContextKey{}).(*verifierContext)
+	r := verifierCtx.r
+	r.Shuffle(len(verifierCtx.randomizedIndexes), func(i, j int) {
+		verifierCtx.randomizedIndexes[i], verifierCtx.randomizedIndexes[j] = verifierCtx.randomizedIndexes[j], verifierCtx.randomizedIndexes[i]
+	})
+}
+
+// DeterministicCompilationVerifierGetRandomizedLocalFunctionIndex returns the randomized index for the given `index`
+// which is assigned by DeterministicCompilationVerifierRandomizeIndexes.
+func DeterministicCompilationVerifierGetRandomizedLocalFunctionIndex(ctx context.Context, index int) int {
+	verifierCtx := ctx.Value(existingValuesContextKey{}).(*verifierContext)
+	ret := verifierCtx.randomizedIndexes[index]
+	return ret
+}
+
+// DeterministicCompilationVerifierSetCurrentFunctionName sets the current function name to the given `functionName`.
+func DeterministicCompilationVerifierSetCurrentFunctionName(ctx context.Context, functionName string) context.Context {
+	return context.WithValue(ctx, currentFunctionNameKey{}, functionName)
+}
+
+// VerifyOrSetDeterministicCompilationContextValue verifies that the `newValue` is the same as the previous value for the given `scope`
+// and the current function name. If the previous value doesn't exist, it sets the value to the given `newValue`.
+//
+// If the verification fails, this prints the diff and exits the process.
+func VerifyOrSetDeterministicCompilationContextValue(ctx context.Context, scope string, newValue string) {
+	fn := ctx.Value(currentFunctionNameKey{}).(string)
+	key := fn + ": " + scope
+	verifierCtx := ctx.Value(existingValuesContextKey{}).(*verifierContext)
+	oldValue, ok := verifierCtx.values[key]
+	if !ok {
+		verifierCtx.values[key] = newValue
+		return
+	}
+	if oldValue != newValue {
+		fmt.Printf(
+			`BUG: Deterministic compilation failed for function%s at scope="%s".
+
+This is mostly due to (but might not be not limited to):
+	* Resetting ssa.Builder, backend.Compiler or frontend.Compiler, etc doens't work as expected, and the compilation has been affected by the previous iterations.
+	* Using a map with non-deterministic iteration order.
+
+---------- [old] ----------
+%s
+
+---------- [new] ----------
+%s
+`,
+			fn,
+			scope, oldValue, newValue,
+		)
+		os.Exit(1)
+	}
+}


Test patch:
diff --git a/internal/engine/wazevo/backend/isa/arm64/util_test.go b/internal/engine/wazevo/backend/isa/arm64/util_test.go
index 1c4cb03a..c9c0f336 100644
--- a/internal/engine/wazevo/backend/isa/arm64/util_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/util_test.go
@@ -1,6 +1,7 @@
 package arm64
 
 import (
+	"context"
 	"strings"
 
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend"
@@ -143,6 +144,6 @@ func (m *mockCompiler) MatchInstrOneOf(def *backend.SSAValueDefinition, opcodes
 }
 
 // Compile implements backend.Compiler.
-func (m *mockCompiler) Compile() (_ []byte, _ []backend.RelocationInfo, goPreambleSize int, _ error) {
+func (m *mockCompiler) Compile(context.Context) (_ []byte, _ []backend.RelocationInfo, goPreambleSize int, _ error) {
 	return
 }