Repository: tetratelabs/wazero
Pull request #169 (2022-01-28 13:25:38)

jit: implement memory.size in assembly.
Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>


Patch:
diff --git a/wasm/jit/engine.go b/wasm/jit/engine.go
index e0cc81f8..99b18199 100644
--- a/wasm/jit/engine.go
+++ b/wasm/jit/engine.go
@@ -443,7 +443,6 @@ func (e *engine) valueStackTopIndex() uint64 {
 
 const (
 	builtinFunctionAddressMemoryGrow wasm.FunctionAddress = iota
-	builtinFunctionAddressMemorySize
 	builtinFunctionAddressGrowValueStack
 	builtinFunctionAddressGrowCallFrameStack
 	// builtinFunctionAddressBreakPoint is internal (only for wazero developers). Disabled by default.
@@ -546,10 +545,6 @@ jitentry:
 			case builtinFunctionAddressMemoryGrow:
 				callerCompiledFunction := e.callFrameTop().compiledFunction
 				e.builtinFunctionMemoryGrow(callerCompiledFunction.source.ModuleInstance.Memory)
-			case builtinFunctionAddressMemorySize:
-				callerCompiledFunction := e.callFrameTop().compiledFunction
-				// TODO: this could be implemented in assembly -- it's just reading engine.memorySliceLen.
-				e.builtinFunctionMemorySize(callerCompiledFunction.source.ModuleInstance.Memory)
 			case builtinFunctionAddressGrowValueStack:
 				callerCompiledFunction := e.callFrameTop().compiledFunction
 				e.builtinFunctionGrowValueStack(callerCompiledFunction.maxStackPointer)
@@ -634,12 +629,14 @@ func (e *engine) builtinFunctionMemoryGrow(mem *wasm.MemoryInstance) {
 	if mem.Max != nil {
 		max = uint64(*mem.Max) * wasm.PageSize
 	}
+
+	currentLen := uint64(len(mem.Buffer))
 	// If exceeds the max of memory size, we push -1 according to the spec.
-	if uint64(newPages*wasm.PageSize+uint64(len(mem.Buffer))) > max {
+	if uint64(newPages*wasm.PageSize+currentLen) > max {
 		v := int32(-1)
 		e.pushValue(uint64(v))
 	} else {
-		e.builtinFunctionMemorySize(mem) // Grow returns the prior memory size on change.
+		e.pushValue(currentLen / wasm.PageSize) // Grow returns the prior memory size on change.
 		mem.Buffer = append(mem.Buffer, make([]byte, newPages*wasm.PageSize)...)
 
 		// Update the moduleContext's fields as they become stale after the update ^^.
@@ -649,10 +646,6 @@ func (e *engine) builtinFunctionMemoryGrow(mem *wasm.MemoryInstance) {
 	}
 }
 
-func (e *engine) builtinFunctionMemorySize(mem *wasm.MemoryInstance) {
-	e.pushValue(uint64(len(mem.Buffer)) / wasm.PageSize)
-}
-
 func (e *engine) Compile(f *wasm.FunctionInstance) (err error) {
 	var compiled *compiledFunction
 	if f.IsHostFunction() {
diff --git a/wasm/jit/jit_amd64.go b/wasm/jit/jit_amd64.go
index ca48d597..0579e766 100644
--- a/wasm/jit/jit_amd64.go
+++ b/wasm/jit/jit_amd64.go
@@ -4282,17 +4282,32 @@ func (c *amd64Compiler) compileMemorySize() error {
 	if err := c.maybeMoveTopConditionalToFreeGeneralPurposeRegister(); err != nil {
 		return err
 	}
-	if err := c.callGoFunction(jitCallStatusCodeCallBuiltInFunction, builtinFunctionAddressMemorySize); err != nil {
+
+	reg, err := c.allocateRegister(generalPurposeRegisterTypeInt)
+	if err != nil {
 		return err
 	}
+	loc := c.locationStack.pushValueOnRegister(reg)
 
-	// After the function call, we have to initialize the stack base pointer and memory reserved registers.
-	c.initializeReservedStackBasePointer()
-	c.initializeReservedMemoryPointer()
-
-	loc := c.locationStack.pushValueOnStack() // The size is pushed on the top.
-	loc.setRegisterType(generalPurposeRegisterTypeInt)
-
+	getMemorySizeInBytes := c.newProg()
+	getMemorySizeInBytes.As = x86.AMOVQ
+	getMemorySizeInBytes.To.Type = obj.TYPE_REG
+	getMemorySizeInBytes.To.Reg = loc.register
+	getMemorySizeInBytes.From.Type = obj.TYPE_MEM
+	getMemorySizeInBytes.From.Reg = reservedRegisterForEngine
+	getMemorySizeInBytes.From.Offset = engineModuleContextMemorySliceLenOffset
+	c.addInstruction(getMemorySizeInBytes)
+
+	// WebAssembly's memory.size returns the page size (65536) of memory region.
+	// That is equivalent to divide the len of memory slice by 65536 and
+	// that can be calculated as SHR by 16 bits as 65536 = 2^16.
+	getMemorySizeInPageUnit := c.newProg()
+	getMemorySizeInPageUnit.As = x86.ASHRQ
+	getMemorySizeInPageUnit.To.Type = obj.TYPE_REG
+	getMemorySizeInPageUnit.To.Reg = loc.register
+	getMemorySizeInPageUnit.From.Type = obj.TYPE_CONST
+	getMemorySizeInPageUnit.From.Offset = 16
+	c.addInstruction(getMemorySizeInPageUnit)
 	return nil
 }
 


Test patch:
diff --git a/wasm/jit/jit_amd64_test.go b/wasm/jit/jit_amd64_test.go
index 2b55d9ae..52893a07 100644
--- a/wasm/jit/jit_amd64_test.go
+++ b/wasm/jit/jit_amd64_test.go
@@ -155,10 +155,12 @@ func (j *jitEnv) execWithModule(code []byte, module *wasm.ModuleInstance) {
 	)
 }
 
+const defaultMemoryPageNumInTest = 2
+
 func newJITEnvironment() *jitEnv {
 	return &jitEnv{
 		eng:   newEngine(),
-		mem:   &wasm.MemoryInstance{Buffer: make([]byte, 1024)},
+		mem:   &wasm.MemoryInstance{Buffer: make([]byte, wasm.PageSize*defaultMemoryPageNumInTest)},
 		table: &wasm.TableInstance{},
 	}
 }
@@ -5466,48 +5468,30 @@ func TestAmd64Compiler_compileMemoryGrow(t *testing.T) {
 }
 
 func TestAmd64Compiler_compileMemorySize(t *testing.T) {
-	for _, currentCallFrameStackPointer := range []uint64{0, 10, 20} {
-		currentCallFrameStackPointer := currentCallFrameStackPointer
-		t.Run(fmt.Sprintf("%d", currentCallFrameStackPointer), func(t *testing.T) {
-			compiler := requireNewCompiler(t)
-			err := compiler.emitPreamble()
-			require.NoError(t, err)
-			// Emit memory.size instructions.
-			err = compiler.compileMemorySize()
-			require.NoError(t, err)
-			// At this point, the size of memory should be pushed onto the stack.
-			require.Equal(t, uint64(1), compiler.locationStack.sp)
-			require.Equal(t, generalPurposeRegisterTypeInt, compiler.locationStack.peek().registerType())
-
-			// Emit arbitrary code after memory.size returned.
-			const expValue uint32 = 100
-			err = compiler.emitPreamble()
-			require.NoError(t, err)
-			err = compiler.compileConstI32(&wazeroir.OperationConstI32{Value: expValue})
-			require.NoError(t, err)
-			err = compiler.releaseAllRegistersToStack()
-			require.NoError(t, err)
-			compiler.exit(jitCallStatusCodeReturned)
-
-			// Generate the code under test.
-			code, _, _, err := compiler.generate()
-			require.NoError(t, err)
+	compiler := requireNewCompiler(t)
+	err := compiler.emitPreamble()
+	require.NoError(t, err)
+	// Emit memory.size instructions.
+	err = compiler.compileMemorySize()
+	require.NoError(t, err)
+	// At this point, the size of memory should be pushed onto the stack.
+	require.Equal(t, uint64(1), compiler.locationStack.sp)
+	require.Equal(t, generalPurposeRegisterTypeInt, compiler.locationStack.peek().registerType())
 
-			// Run code.
-			env := newJITEnvironment()
-			env.exec(code)
+	err = compiler.releaseAllRegistersToStack()
+	require.NoError(t, err)
+	compiler.exit(jitCallStatusCodeReturned)
 
-			require.Equal(t, jitCallStatusCodeCallBuiltInFunction, env.jitStatus())
-			require.Equal(t, wasm.FunctionAddress(builtinFunctionAddressMemorySize), env.functionCallAddress())
+	// Generate the code under test.
+	code, _, _, err := compiler.generate()
+	require.NoError(t, err)
 
-			returnAddress := env.callFrameStackPeek().returnAddress
-			require.NotZero(t, returnAddress)
-			jitcall(returnAddress, uintptr(unsafe.Pointer(env.engine())))
+	// Run code.
+	env := newJITEnvironment()
+	env.exec(code)
 
-			require.Equal(t, expValue, env.stackTopAsUint32())
-			require.Equal(t, jitCallStatusCodeReturned, env.jitStatus())
-		})
-	}
+	require.Equal(t, jitCallStatusCodeReturned, env.jitStatus())
+	require.Equal(t, uint32(defaultMemoryPageNumInTest), env.stackTopAsUint32())
 }
 
 func TestAmd64Compiler_compileDrop(t *testing.T) {