Repository: tetratelabs/wazero
Pull request #68 (2021-12-09 12:19:55)

jit: grow stack automatically. 
This commit resolves one TODO for growing stack automatically 
when the engine enters new JITed function frames and sees the necessity
to grow the stack.

Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>

part of #65.

mathetake: I will put the comments on the struct and then merge this :D

Patch:
diff --git a/wasm/jit/engine.go b/wasm/jit/engine.go
index eb628387..1dd9820a 100644
--- a/wasm/jit/engine.go
+++ b/wasm/jit/engine.go
@@ -160,7 +160,7 @@ func NewEngine() wasm.Engine {
 	return newEngine()
 }
 
-const initialStackSize = 100
+const initialStackSize = 1024
 
 func newEngine() *engine {
 	e := &engine{
@@ -251,16 +251,30 @@ type compiledWasmFunction struct {
 	codeInitialAddress uintptr
 	// The same purpose as codeInitialAddress, but for memory.Buffer.
 	memoryAddress uintptr
+	// The max of the stack pointer this function can reach. Lazily applied via maybeGrowStack.
+	maxStackPointer uint64
 }
 
 const (
 	builtinFunctionIndexGrowMemory = iota
 )
 
-func (e *engine) stackGrow() {
-	newStack := make([]uint64, len(e.stack)*2)
-	copy(newStack[:len(e.stack)], e.stack)
-	e.stack = newStack
+// Grow the stack size according to maxStackPointer argument
+// which is the max stack pointer from the base pointer
+// for the next function frame execution.
+func (e *engine) maybeGrowStack(maxStackPointer uint64) {
+	currentLen := uint64(len(e.stack))
+	remained := currentLen - e.currentBaseStackPointer
+	if maxStackPointer > remained {
+		// This case we need to grow the stack as the empty slots
+		// are not able to store all the stack items.
+		// So we grow the stack with the new len = currentLen*2+maxStackPointer.
+		newStack := make([]uint64, currentLen*2+(maxStackPointer))
+		top := e.currentBaseStackPointer + e.currentStackPointer
+		copy(newStack[:top], e.stack[:top])
+		e.stack = newStack
+	}
+	// TODO: maybe better think about how to shrink the stack as well.
 }
 
 func (e *engine) exec(f *compiledWasmFunction) {
@@ -270,13 +284,8 @@ func (e *engine) exec(f *compiledWasmFunction) {
 		caller:                   nil,
 		continuationStackPointer: f.inputs,
 	}
-	// TODO: We should check the size of the stack,
-	// and if it's running out, grow it before calling into JITed code.
-	// It should be possible to check the necessity by statically
-	// analyzing the max height of the stack in the function.
-	if false {
-		e.stackGrow()
-	}
+	// If the Go-allocated stack is running out, we grow it before calling into JITed code.
+	e.maybeGrowStack(f.maxStackPointer)
 	for e.callFrameStack != nil {
 		currentFrame := e.callFrameStack
 		if buildoptions.IsDebugMode {
@@ -322,13 +331,8 @@ func (e *engine) exec(f *compiledWasmFunction) {
 				// Set the base pointer to the beginning of the function inputs
 				baseStackPointer: e.currentBaseStackPointer + e.currentStackPointer - nextFunc.inputs,
 			}
-			// TODO: We should check the size of the stack,
-			// and if it's running out, grow it before calling into JITed code.
-			// It should be possible to check the necessity by statically
-			// analyzing the max height of the stack in the function.
-			if false {
-				e.stackGrow()
-			}
+			// If the Go-allocated stack is running out, we grow it before calling into JITed code.
+			e.maybeGrowStack(nextFunc.maxStackPointer)
 			// Now move onto the callee function.
 			e.callFrameStack = frame
 			e.currentBaseStackPointer = frame.baseStackPointer
diff --git a/wasm/jit/jit_amd64.go b/wasm/jit/jit_amd64.go
index c746e4e6..5c9b7024 100644
--- a/wasm/jit/jit_amd64.go
+++ b/wasm/jit/jit_amd64.go
@@ -224,10 +224,11 @@ func (e *engine) compileWasmFunction(f *wasm.FunctionInstance) (*compiledWasmFun
 
 func (b *amd64Builder) newCompiledWasmFunction(code []byte) *compiledWasmFunction {
 	cf := &compiledWasmFunction{
-		codeSegment: code,
-		inputs:      uint64(len(b.f.Signature.InputTypes)),
-		returns:     uint64(len(b.f.Signature.ReturnTypes)),
-		memory:      b.f.ModuleInstance.Memory,
+		codeSegment:     code,
+		inputs:          uint64(len(b.f.Signature.InputTypes)),
+		returns:         uint64(len(b.f.Signature.ReturnTypes)),
+		memory:          b.f.ModuleInstance.Memory,
+		maxStackPointer: b.locationStack.maxStackPointer,
 	}
 	if cf.memory != nil {
 		cf.memoryAddress = uintptr(unsafe.Pointer(&cf.memory.Buffer[0]))
diff --git a/wasm/jit/jit_value_location_amd64.go b/wasm/jit/jit_value_location_amd64.go
index f0d22f6a..aeef7b24 100644
--- a/wasm/jit/jit_value_location_amd64.go
+++ b/wasm/jit/jit_value_location_amd64.go
@@ -54,6 +54,9 @@ const (
 	conditionalRegisterStateBE                             // CF | ZF below or equal (unsigned <=)
 )
 
+// valueLocation corresponds to each variable pushed onto the wazeroir (virtual) stack,
+// and it has the information about where it exists in the physical machine.
+// It might exist in registers, or maybe on in the non-virtual physical stack allocated in memory.
 type valueLocation struct {
 	valueType wazeroir.SignLessType
 	// Set to -1 if the value is stored in the memory stack.
@@ -76,6 +79,7 @@ func (v *valueLocation) registerType() (t generalPurposeRegisterType) {
 	}
 	return
 }
+
 func (v *valueLocation) setValueType(t wazeroir.SignLessType) {
 	v.valueType = t
 }
@@ -103,10 +107,24 @@ func newValueLocationStack() *valueLocationStack {
 	}
 }
 
+// valueLocationStack represents the wazeroir virtual stack
+// where each item holds the information about where it exists
+// on the physical machine.
+// Notably this is only used in the compilation phase, not runtime,
+// and we change the state of this struct at every wazeroir operation we compile.
+// In this way, we can see where the operands of a operation (for example,
+// two variables for wazeroir add operation.) exist and check the neccesity for
+// moving the variable to registers to perform actual CPU instruction
+// to achieve wazeroir's add operation.
 type valueLocationStack struct {
-	stack         []*valueLocation
-	sp            uint64
+	// Holds all the variables.
+	stack []*valueLocation
+	// The current stack pointer.
+	sp uint64
+	// Stores the used registers.
 	usedRegisters map[int16]struct{}
+	// Records max(.sp) across the lifespan of this struct.
+	maxStackPointer uint64
 }
 
 func (s *valueLocationStack) clone() *valueLocationStack {
@@ -127,6 +145,7 @@ func (s *valueLocationStack) clone() *valueLocationStack {
 	}
 	return ret
 }
+
 func (s *valueLocationStack) pushValueOnRegister(reg int16) (loc *valueLocation) {
 	loc = &valueLocation{register: reg, conditionalRegister: conditionalRegisterStateUnset}
 	s.markRegisterUsed(reg)
@@ -155,6 +174,9 @@ func (s *valueLocationStack) push(loc *valueLocation) {
 	} else {
 		s.stack[s.sp] = loc
 	}
+	if s.sp > s.maxStackPointer {
+		s.maxStackPointer = s.sp
+	}
 	s.sp++
 }
 


Test patch:
diff --git a/wasm/jit/engine_test.go b/wasm/jit/engine_test.go
index 540ed723..63007666 100644
--- a/wasm/jit/engine_test.go
+++ b/wasm/jit/engine_test.go
@@ -85,16 +85,35 @@ func TestEngine_PreCompile(t *testing.T) {
 	require.Equal(t, prevCompiledFunctions, eng.compiledWasmFunctions)
 }
 
-func TestEngine_stackGrow(t *testing.T) {
-	eng := newEngine()
-	require.Len(t, eng.stack, initialStackSize)
-	eng.push(10)
-	require.Equal(t, uint64(1), eng.currentStackPointer)
-	require.Equal(t, uint64(10), eng.stack[eng.currentStackPointer-1])
-	eng.stackGrow()
-	require.Len(t, eng.stack, initialStackSize*2)
-	// stackGrow only grows the stack len,
-	// and must not modify neither stack pointer nor the values in the stack.
-	require.Equal(t, uint64(1), eng.currentStackPointer)
-	require.Equal(t, uint64(10), eng.stack[eng.currentStackPointer-1])
+func TestEngine_maybeGrowStack(t *testing.T) {
+	t.Run("grow", func(t *testing.T) {
+		eng := &engine{stack: make([]uint64, 10)}
+		eng.currentBaseStackPointer = 5
+		eng.push(10)
+		require.Equal(t, uint64(1), eng.currentStackPointer)
+		require.Equal(t, uint64(10), eng.stack[eng.currentBaseStackPointer+eng.currentStackPointer-1])
+		eng.maybeGrowStack(100)
+		// Currently we have 9 empty slots (10 - 1(base pointer)) above base pointer for new items,
+		// but we require 100 max stack pointer for the next function,
+		// so this results in making the stack length 120 = 10(current len)*2+(100(maxStackPointer))
+		require.Len(t, eng.stack, 120)
+		// maybeAdjustStack only shrink the stack,
+		// and must not modify neither stack pointer nor the values in the stack.
+		require.Equal(t, uint64(1), eng.currentStackPointer)
+		require.Equal(t, uint64(10), eng.stack[eng.currentBaseStackPointer+eng.currentStackPointer-1])
+	})
+	t.Run("noop", func(t *testing.T) {
+		eng := &engine{stack: make([]uint64, 10)}
+		eng.currentBaseStackPointer = 1
+		eng.push(10)
+		require.Equal(t, uint64(1), eng.currentStackPointer)
+		require.Equal(t, uint64(10), eng.stack[eng.currentBaseStackPointer+eng.currentStackPointer-1])
+		eng.maybeGrowStack(6)
+		// Currently we have 9 empty slots (10 - 1(base pointer)) above base pointer for new items,
+		// and we only require 6 max stack pointer for the next function, so we have enough empty slots.
+		// so maybeGrowStack must not modify neither stack pointer, the values in the stack nor stack len.
+		require.Len(t, eng.stack, 10)
+		require.Equal(t, uint64(1), eng.currentStackPointer)
+		require.Equal(t, uint64(10), eng.stack[eng.currentBaseStackPointer+eng.currentStackPointer-1])
+	})
 }
diff --git a/wasm/jit/jit_value_location_amd64_test.go b/wasm/jit/jit_value_location_amd64_test.go
index 2d9afb7f..01eacf6f 100644
--- a/wasm/jit/jit_value_location_amd64_test.go
+++ b/wasm/jit/jit_value_location_amd64_test.go
@@ -51,6 +51,14 @@ func TestValueLocationStack_basic(t *testing.T) {
 		actual, exp := s.stack[i], cloned.stack[i]
 		require.NotEqual(t, uintptr(unsafe.Pointer(exp)), uintptr(unsafe.Pointer(actual)))
 	}
+	// Check the max stack pointer.
+	for i := 0; i < 1000; i++ {
+		s.pushValueOnStack()
+	}
+	for i := 0; i < 1000; i++ {
+		s.pop()
+	}
+	require.Equal(t, uint64(1001), s.maxStackPointer)
 }
 
 func TestValueLocationStack_takeFreeRegister(t *testing.T) {