Repository: tetratelabs/wazero
Pull request #728 (2022-07-30 02:07:32)

compiler: support for Wasm-implemented host functinos
Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>


Patch:
diff --git a/internal/engine/compiler/compiler.go b/internal/engine/compiler/compiler.go
index b5661fdc..079c604e 100644
--- a/internal/engine/compiler/compiler.go
+++ b/internal/engine/compiler/compiler.go
@@ -16,9 +16,9 @@ type compiler interface {
 	// compile generates the byte slice of native code.
 	// stackPointerCeil is the max stack pointer that the target function would reach.
 	compile() (code []byte, stackPointerCeil uint64, err error)
-	// compileHostFunction adds the trampoline code from which native code can jump into the host function.
+	// compileGoHostFunction adds the trampoline code from which native code can jump into the Go-defined host function.
 	// TODO: maybe we wouldn't need to have trampoline for host functions.
-	compileHostFunction() error
+	compileGoDefinedHostFunction() error
 	// compileLabel notify compilers of the beginning of a label.
 	// Return true if the compiler decided to skip the entire label.
 	// See wazeroir.OperationLabel
diff --git a/internal/engine/compiler/engine.go b/internal/engine/compiler/engine.go
index ea245f48..a73ab924 100644
--- a/internal/engine/compiler/engine.go
+++ b/internal/engine/compiler/engine.go
@@ -424,13 +424,10 @@ func (e *engine) CompileModule(ctx context.Context, module *wasm.Module) error {
 	for funcIndex, ir := range irs {
 		var compiled *code
 		if ir.GoFunc != nil {
-			if compiled, err = compileHostFunction(ir); err != nil {
+			if compiled, err = compileGoDefinedHostFunction(ir); err != nil {
 				def := module.FunctionDefinitionSection[uint32(funcIndex)+module.ImportFuncCount()]
 				return fmt.Errorf("error compiling host go func[%s]: %w", def.DebugName(), err)
 			}
-		} else if ir.IsHostFunction && ir.UsesMemory {
-			def := module.FunctionDefinitionSection[uint32(funcIndex)+module.ImportFuncCount()]
-			return fmt.Errorf("error compiling host wasm func[%s]: memory access not yet supported", def.DebugName())
 		} else if compiled, err = compileWasmFunction(e.enabledFeatures, ir); err != nil {
 			def := module.FunctionDefinitionSection[uint32(funcIndex)+module.ImportFuncCount()]
 			return fmt.Errorf("error compiling wasm func[%s]: %w", def.DebugName(), err)
@@ -807,14 +804,13 @@ func (ce *callEngine) builtinFunctionTableGrow(ctx context.Context, tables []*wa
 	ce.pushValue(uint64(res))
 }
 
-func compileHostFunction(ir *wazeroir.CompilationResult) (*code, error) {
+func compileGoDefinedHostFunction(ir *wazeroir.CompilationResult) (*code, error) {
 	compiler, err := newCompiler(ir)
 	if err != nil {
 		return nil, err
 	}
 
-	// TODO: Consider when no memory is used (ex !ir.UsesMemory)
-	if err = compiler.compileHostFunction(); err != nil {
+	if err = compiler.compileGoDefinedHostFunction(); err != nil {
 		return nil, err
 	}
 
diff --git a/internal/engine/compiler/impl_amd64.go b/internal/engine/compiler/impl_amd64.go
index 54ce99a7..bfb4dcea 100644
--- a/internal/engine/compiler/impl_amd64.go
+++ b/internal/engine/compiler/impl_amd64.go
@@ -156,9 +156,9 @@ func (c *amd64Compiler) label(labelKey string) *amd64LabelInfo {
 	return c.labels[labelKey]
 }
 
-// compileHostFunction constructs the entire code to enter the host function implementation,
-// and return back to the caller.
-func (c *amd64Compiler) compileHostFunction() error {
+// compileGoDefinedHostFunction constructs the entire code to enter the host function implementation,
+// and return to the caller.
+func (c *amd64Compiler) compileGoDefinedHostFunction() error {
 	// First we must update the location stack to reflect the number of host function inputs.
 	c.pushFunctionParams()
 
@@ -4934,6 +4934,13 @@ func (c *amd64Compiler) compileMaybeGrowValueStack() error {
 // callEngine.ModuleContext.ModuleInstanceAddress.
 // This is called in two cases: in function preamble, and on the return from (non-Go) function calls.
 func (c *amd64Compiler) compileModuleContextInitialization() error {
+	if c.ir.IsHostFunction {
+		// If this is the host function, our semantic is that host functions use the caller's context.
+		// Therefore, we can skip all the initialization steps here as at this point, the caller's memory, etc
+		// are already set on the callEngine.
+		return nil
+	}
+
 	// amd64CallingConventionModuleInstanceAddressRegister holds the module instance's address
 	// so mark it used so that it won't be used as a free register until the module context initialization finishes.
 	c.locationStack.markRegisterUsed(amd64CallingConventionModuleInstanceAddressRegister)
diff --git a/internal/engine/compiler/impl_arm64.go b/internal/engine/compiler/impl_arm64.go
index 9fedea3f..f370174a 100644
--- a/internal/engine/compiler/impl_arm64.go
+++ b/internal/engine/compiler/impl_arm64.go
@@ -398,8 +398,8 @@ func (c *arm64Compiler) compileExitFromNativeCode(status nativeCallStatusCode) {
 	c.assembler.CompileJumpToRegister(arm64.RET, arm64ReservedRegisterForTemporary)
 }
 
-// compileHostFunction implements compiler.compileHostFunction for the arm64 architecture.
-func (c *arm64Compiler) compileHostFunction() error {
+// compileGoHostFunction implements compiler.compileHostFunction for the arm64 architecture.
+func (c *arm64Compiler) compileGoDefinedHostFunction() error {
 	// First we must update the location stack to reflect the number of host function inputs.
 	c.pushFunctionParams()
 
@@ -4131,6 +4131,13 @@ func (c *arm64Compiler) compileReservedMemoryRegisterInitialization() {
 // ce.moduleContext.ModuleInstanceAddress.
 // This is called in two cases: in function preamble, and on the return from (non-Go) function calls.
 func (c *arm64Compiler) compileModuleContextInitialization() error {
+	if c.ir.IsHostFunction {
+		// If this is the host function, our semantic is that host functions use the caller's context.
+		// Therefore, we can skip all the initialization steps here as at this point, the caller's memory, etc
+		// are already set on the callEngine.
+		return nil
+	}
+
 	c.markRegisterUsed(arm64CallingConventionModuleInstanceAddressRegister)
 	defer c.markRegisterUnused(arm64CallingConventionModuleInstanceAddressRegister)
 


Test patch:
diff --git a/internal/engine/compiler/compiler_controlflow_test.go b/internal/engine/compiler/compiler_controlflow_test.go
index 94596a6d..f1642831 100644
--- a/internal/engine/compiler/compiler_controlflow_test.go
+++ b/internal/engine/compiler/compiler_controlflow_test.go
@@ -14,7 +14,7 @@ func TestCompiler_compileHostFunction(t *testing.T) {
 	env := newCompilerEnvironment()
 	compiler := env.requireNewCompiler(t, newCompiler, nil)
 
-	err := compiler.compileHostFunction()
+	err := compiler.compileGoDefinedHostFunction()
 	require.NoError(t, err)
 
 	// Generate and run the code under test.
diff --git a/internal/testing/enginetest/enginetest.go b/internal/testing/enginetest/enginetest.go
index 531d1546..08283031 100644
--- a/internal/testing/enginetest/enginetest.go
+++ b/internal/testing/enginetest/enginetest.go
@@ -374,9 +374,7 @@ func runTestModuleEngine_Call_HostFn_Mem(t *testing.T, et EngineTester, readMem
 func RunTestModuleEngine_Call_HostFn(t *testing.T, et EngineTester) {
 	t.Run("wasm", func(t *testing.T) {
 		runTestModuleEngine_Call_HostFn(t, et, hostDivByWasm)
-		if !et.IsCompiler() { // TODO: Support host wasm func that uses caller memory in compiler.
-			runTestModuleEngine_Call_HostFn_Mem(t, et, hostReadMemWasm)
-		}
+		runTestModuleEngine_Call_HostFn_Mem(t, et, hostReadMemWasm)
 	})
 	t.Run("go", func(t *testing.T) {
 		runTestModuleEngine_Call_HostFn(t, et, hostDivByGo)
@@ -844,7 +842,9 @@ func setupCallMemTests(t *testing.T, e wasm.Engine, readMem *wasm.Code, fnlf exp
 				{Index: 3, Name: callImportCallReadMemName},
 			},
 		},
-		ID: wasm.ModuleID{1},
+		// Indicates that this module has a memory so that compilers are able to assembe memory-related initialization.
+		MemorySection: &wasm.Memory{Min: 1},
+		ID:            wasm.ModuleID{1},
 	}
 	importingModule.BuildFunctionDefinitions()
 	err = e.CompileModule(testCtx, importingModule)