Repository: tetratelabs/wazero
Pull request #81 (2021-12-22 02:40:59)

jit: support global get and set instructions
part of #65 


Patch:
diff --git a/wasm/jit/engine.go b/wasm/jit/engine.go
index bad7e388..822e960d 100644
--- a/wasm/jit/engine.go
+++ b/wasm/jit/engine.go
@@ -12,6 +12,8 @@ import (
 )
 
 type engine struct {
+	// These fields are used and manipulated by JITed code.
+
 	// The actual Go-allocated stack.
 	// Note that we NEVER edit len or cap in JITed code so we won't get screwed when GC comes in.
 	stack []uint64
@@ -29,8 +31,13 @@ type engine struct {
 	// Instructions after [base+continuationAddressOffset] must start with
 	// restoring reserved registeres.
 	continuationAddressOffset uintptr
+	// The current compiledWasmFunction.globalSliceAddress
+	currentGlobalSliceAddress uintptr
 	// Function call frames in linked list
 	callFrameStack *callFrame
+
+	// The following fields are only used during compilation.
+
 	// Store the compiled functions and indexes.
 	compiledWasmFunctions     []*compiledWasmFunction
 	compiledWasmFunctionIndex map[*wasm.FunctionInstance]int64
@@ -39,6 +46,17 @@ type engine struct {
 	compiledHostFunctionIndex map[*wasm.FunctionInstance]int64
 }
 
+// Native code manipulates the engine's fields with these constants.
+const (
+	engineStackSliceOffset                = 0
+	engineCurrentStackPointerOffset       = 24
+	engineCurrentBaseStackPointerOffset   = 32
+	engineJITCallStatusCodeOffset         = 40
+	engineFunctionCallIndexOffset         = 48
+	engineContinuationAddressOffset       = 56
+	engineCurrentGlobalSliceAddressOffset = 64
+)
+
 func (e *engine) Call(f *wasm.FunctionInstance, args ...uint64) (returns []uint64, err error) {
 	prevFrame := e.callFrameStack
 	// We ensure that this Call method never panics as
@@ -244,16 +262,6 @@ func (s jitCallStatusCode) String() (ret string) {
 	return
 }
 
-// These consts are used in native codes to manipulate the engine's fields.
-const (
-	engineStackSliceOffset              = 0
-	engineCurrentStackPointerOffset     = 24
-	engineCurrentBaseStackPointerOffset = 32
-	engineJITCallStatusCodeOffset       = 40
-	engineFunctionCallIndexOffset       = 48
-	engineContinuationAddressOffset     = 56
-)
-
 type callFrame struct {
 	continuationAddress      uintptr
 	continuationStackPointer uint64
@@ -299,6 +307,8 @@ type compiledWasmFunction struct {
 	codeInitialAddress uintptr
 	// The same purpose as codeInitialAddress, but for memory.Buffer.
 	memoryAddress uintptr
+	// globalSliceAddress is like codeInitialAddress, but for .globals.
+	globalSliceAddress uintptr
 	// The max of the stack pointer this function can reach. Lazily applied via maybeGrowStack.
 	maxStackPointer uint64
 }
@@ -332,6 +342,7 @@ func (e *engine) exec(f *compiledWasmFunction) {
 		caller:                   nil,
 		continuationStackPointer: f.inputs,
 	}
+	e.currentGlobalSliceAddress = f.globalSliceAddress
 	// If the Go-allocated stack is running out, we grow it before calling into JITed code.
 	e.maybeGrowStack(f.maxStackPointer)
 	for e.callFrameStack != nil {
@@ -386,6 +397,7 @@ func (e *engine) exec(f *compiledWasmFunction) {
 			e.currentBaseStackPointer = frame.baseStackPointer
 			// Set the stack pointer so that base+sp would point to the top of function inputs.
 			e.currentStackPointer = nextFunc.inputs
+			e.currentGlobalSliceAddress = nextFunc.globalSliceAddress
 		case jitCallStatusCodeCallBuiltInFunction:
 			switch e.functionCallIndex {
 			case builtinFunctionIndexGrowMemory:
@@ -399,7 +411,7 @@ func (e *engine) exec(f *compiledWasmFunction) {
 			// Push the call frame for this host function.
 			e.callFrameStack = &callFrame{hostFunction: targetHostFunction, caller: currentFrame}
 			// Call into the host function.
-			targetHostFunction.f(&wasm.HostFunctionCallContext{Memory: f.memory})
+			targetHostFunction.f(&wasm.HostFunctionCallContext{Memory: currentFrame.wasmFunction.memory})
 			// Pop the call frame.
 			e.callFrameStack = currentFrame
 		case jitCallStatusCodeUnreachable:
diff --git a/wasm/jit/jit_amd64.go b/wasm/jit/jit_amd64.go
index bd9ff4dc..a4e2983a 100644
--- a/wasm/jit/jit_amd64.go
+++ b/wasm/jit/jit_amd64.go
@@ -96,9 +96,13 @@ func (e *engine) compileWasmFunction(f *wasm.FunctionInstance) (*compiledWasmFun
 				return nil, fmt.Errorf("error handling swap operation: %w", err)
 			}
 		case *wazeroir.OperationGlobalGet:
-			return nil, fmt.Errorf("unsupported operation in JIT compiler: %v", o)
+			if err := builder.handleGlobalGet(o); err != nil {
+				return nil, fmt.Errorf("error handling global.get operation: %w", err)
+			}
 		case *wazeroir.OperationGlobalSet:
-			return nil, fmt.Errorf("unsupported operation in JIT compiler: %v", o)
+			if err := builder.handleGlobalSet(o); err != nil {
+				return nil, fmt.Errorf("error handling global.set operation: %w", err)
+			}
 		case *wazeroir.OperationLoad:
 			return nil, fmt.Errorf("unsupported operation in JIT compiler: %v", o)
 		case *wazeroir.OperationLoad8:
@@ -239,6 +243,9 @@ func (b *amd64Builder) newCompiledWasmFunction(code []byte) *compiledWasmFunctio
 	if cf.memory != nil {
 		cf.memoryAddress = uintptr(unsafe.Pointer(&cf.memory.Buffer[0]))
 	}
+	if len(b.f.ModuleInstance.Globals) > 0 {
+		cf.globalSliceAddress = uintptr(unsafe.Pointer(&b.f.ModuleInstance.Globals[0]))
+	}
 	cf.codeInitialAddress = uintptr(unsafe.Pointer(&cf.codeSegment[0]))
 	return cf
 }
@@ -395,6 +402,131 @@ func (b *amd64Builder) handleSwap(o *wazeroir.OperationSwap) error {
 	return nil
 }
 
+const globalInstanceValueOffset = 8
+
+func (b *amd64Builder) handleGlobalGet(o *wazeroir.OperationGlobalGet) error {
+	intReg, err := b.allocateRegister(gpTypeInt)
+	if err != nil {
+		return err
+	}
+
+	// First, move the pointer to the global slice into the allocated register.
+	moveGlobalSlicePointer := b.newProg()
+	moveGlobalSlicePointer.As = x86.AMOVQ
+	moveGlobalSlicePointer.To.Type = obj.TYPE_REG
+	moveGlobalSlicePointer.To.Reg = intReg
+	moveGlobalSlicePointer.From.Type = obj.TYPE_MEM
+	moveGlobalSlicePointer.From.Reg = engineInstanceReg
+	moveGlobalSlicePointer.From.Offset = engineCurrentGlobalSliceAddressOffset
+	b.addInstruction(moveGlobalSlicePointer)
+
+	// Then, get the memory location of the target global instance's pointer.
+	getGlobalInstanceLocation := b.newProg()
+	getGlobalInstanceLocation.As = x86.AADDQ
+	getGlobalInstanceLocation.To.Type = obj.TYPE_REG
+	getGlobalInstanceLocation.To.Reg = intReg
+	getGlobalInstanceLocation.From.Type = obj.TYPE_CONST
+	getGlobalInstanceLocation.From.Offset = 8 * int64(o.Index)
+	b.addInstruction(getGlobalInstanceLocation)
+
+	// Now, move the location of the global instance into the register.
+	getGlobalInstancePointer := b.newProg()
+	getGlobalInstancePointer.As = x86.AMOVQ
+	getGlobalInstancePointer.To.Type = obj.TYPE_REG
+	getGlobalInstancePointer.To.Reg = intReg
+	getGlobalInstancePointer.From.Type = obj.TYPE_MEM
+	getGlobalInstancePointer.From.Reg = intReg
+	b.addInstruction(getGlobalInstancePointer)
+
+	// When an integer, reuse the pointer register for the value. Otherwise, allocate a float register for it.
+	valueReg := intReg
+	wasmType := b.f.ModuleInstance.Globals[o.Index].Type.ValType
+	switch wasmType {
+	case wasm.ValueTypeF32, wasm.ValueTypeF64:
+		valueReg, err = b.allocateRegister(gpTypeFloat)
+		if err != nil {
+			return err
+		}
+	}
+
+	// Using the register holding the pointer to the target instance, move its value into a register.
+	moveValue := b.newProg()
+	moveValue.As = x86.AMOVQ
+	moveValue.To.Type = obj.TYPE_REG
+	moveValue.To.Reg = valueReg
+	moveValue.From.Type = obj.TYPE_MEM
+	moveValue.From.Reg = intReg
+	moveValue.From.Offset = globalInstanceValueOffset
+	b.addInstruction(moveValue)
+
+	// Record that the retrieved global value on the top of the stack is now in a register.
+	loc := b.locationStack.pushValueOnRegister(valueReg)
+	loc.setValueType(wazeroir.WasmValueTypeToSignless(wasmType))
+	return nil
+}
+
+func (b *amd64Builder) handleGlobalSet(o *wazeroir.OperationGlobalSet) error {
+	// First, move the value to set into a temporary register.
+	val := b.locationStack.pop()
+	if val.onStack() {
+		if err := b.moveStackToRegisterWithAllocation(val.registerType(), val); err != nil {
+			return err
+		}
+	} else if val.onConditionalRegister() {
+		if err := b.moveConditionalToGPRegister(val); err != nil {
+			return err
+		}
+	}
+
+	// Allocate a register to hold the memory location of the target global instance.
+	intReg, err := b.allocateRegister(gpTypeInt)
+	if err != nil {
+		return err
+	}
+
+	// First, move the pointer to the global slice into the allocated register.
+	moveGlobalSlicePointer := b.newProg()
+	moveGlobalSlicePointer.As = x86.AMOVQ
+	moveGlobalSlicePointer.To.Type = obj.TYPE_REG
+	moveGlobalSlicePointer.To.Reg = intReg
+	moveGlobalSlicePointer.From.Type = obj.TYPE_MEM
+	moveGlobalSlicePointer.From.Reg = engineInstanceReg
+	moveGlobalSlicePointer.From.Offset = engineCurrentGlobalSliceAddressOffset
+	b.addInstruction(moveGlobalSlicePointer)
+
+	// Then, get the memory location of the target global instance's pointer.
+	getGlobalInstanceLocation := b.newProg()
+	getGlobalInstanceLocation.As = x86.AADDQ
+	getGlobalInstanceLocation.To.Type = obj.TYPE_REG
+	getGlobalInstanceLocation.To.Reg = intReg
+	getGlobalInstanceLocation.From.Type = obj.TYPE_CONST
+	getGlobalInstanceLocation.From.Offset = 8 * int64(o.Index)
+	b.addInstruction(getGlobalInstanceLocation)
+
+	// Now, move the location of the global instance into the register.
+	getGlobalInstancePointer := b.newProg()
+	getGlobalInstancePointer.As = x86.AMOVQ
+	getGlobalInstancePointer.To.Type = obj.TYPE_REG
+	getGlobalInstancePointer.To.Reg = intReg
+	getGlobalInstancePointer.From.Type = obj.TYPE_MEM
+	getGlobalInstancePointer.From.Reg = intReg
+	b.addInstruction(getGlobalInstancePointer)
+
+	// Now ready to write the value to the global instance location.
+	moveValue := b.newProg()
+	moveValue.As = x86.AMOVQ
+	moveValue.From.Type = obj.TYPE_REG
+	moveValue.From.Reg = val.register
+	moveValue.To.Type = obj.TYPE_MEM
+	moveValue.To.Reg = intReg
+	moveValue.To.Offset = globalInstanceValueOffset
+	b.addInstruction(moveValue)
+
+	// Since the value is now written to memory, release the value register.
+	b.locationStack.releaseRegister(val)
+	return nil
+}
+
 func (b *amd64Builder) handleBr(o *wazeroir.OperationBr) error {
 	if o.Target.IsReturnTarget() {
 		// Release all the registers as our calling convention requires the callee-save.


Test patch:
diff --git a/wasm/jit/engine_test.go b/wasm/jit/engine_test.go
index b434c4f8..7f62c9e1 100644
--- a/wasm/jit/engine_test.go
+++ b/wasm/jit/engine_test.go
@@ -20,6 +20,7 @@ func TestEngine_veifyOffsetValue(t *testing.T) {
 	require.Equal(t, int(unsafe.Offsetof((&engine{}).jitCallStatusCode)), engineJITCallStatusCodeOffset)
 	require.Equal(t, int(unsafe.Offsetof((&engine{}).functionCallIndex)), engineFunctionCallIndexOffset)
 	require.Equal(t, int(unsafe.Offsetof((&engine{}).continuationAddressOffset)), engineContinuationAddressOffset)
+	require.Equal(t, int(unsafe.Offsetof((&engine{}).currentGlobalSliceAddress)), engineCurrentGlobalSliceAddressOffset)
 }
 
 func TestEngine_fibonacci(t *testing.T) {
diff --git a/wasm/jit/jit_amd64_test.go b/wasm/jit/jit_amd64_test.go
index c32a8e14..0e55b074 100644
--- a/wasm/jit/jit_amd64_test.go
+++ b/wasm/jit/jit_amd64_test.go
@@ -1682,3 +1682,99 @@ func TestAmd64Builder_handleSwap(t *testing.T) {
 		})
 	}
 }
+
+// TestGlobalInstanceValueOffset ensures the globalInstanceValueOffset doesn't drift when we modify the struct (wasm.GlobalInstance).
+func TestGlobalInstanceValueOffset(t *testing.T) {
+	require.Equal(t, int(unsafe.Offsetof((&wasm.GlobalInstance{}).Val)), globalInstanceValueOffset)
+}
+
+func TestAmd64Builder_handleGlobalGet(t *testing.T) {
+	const globalValue uint64 = 12345
+	for i, tp := range []wasm.ValueType{
+		wasm.ValueTypeF32, wasm.ValueTypeF64, wasm.ValueTypeI32, wasm.ValueTypeI64,
+	} {
+		tp := tp
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+			// Setup the globals.
+			builder := requireNewBuilder(t)
+			globals := []*wasm.GlobalInstance{nil, {Val: globalValue, Type: &wasm.GlobalType{ValType: tp}}, nil}
+			builder.f = &wasm.FunctionInstance{ModuleInstance: &wasm.ModuleInstance{Globals: globals}}
+			// Emit the code.
+			builder.initializeReservedRegisters()
+			op := &wazeroir.OperationGlobalGet{Index: 1}
+			err := builder.handleGlobalGet(op)
+			require.NoError(t, err)
+
+			// At this point, the top of stack must be the retrieved global on a register.
+			global := builder.locationStack.peek()
+			require.True(t, global.onRegister())
+			require.Len(t, builder.locationStack.usedRegisters, 1)
+			switch tp {
+			case wasm.ValueTypeF32, wasm.ValueTypeF64:
+				require.True(t, isFloatRegister(global.register))
+			case wasm.ValueTypeI32, wasm.ValueTypeI64:
+				require.True(t, isIntRegister(global.register))
+			}
+			builder.releaseAllRegistersToStack()
+			builder.returnFunction()
+
+			// Assemble.
+			code, err := builder.assemble()
+			require.NoError(t, err)
+
+			// Run the code assembled above.
+			eng := newEngine()
+			eng.currentGlobalSliceAddress = uintptr(unsafe.Pointer(&globals[0]))
+			mem := newMemoryInst()
+			jitcall(
+				uintptr(unsafe.Pointer(&code[0])),
+				uintptr(unsafe.Pointer(eng)),
+				uintptr(unsafe.Pointer(&mem.Buffer[0])),
+			)
+			// Since we call global.get, the top of the stack must be the global value.
+			require.Equal(t, globalValue, eng.stack[0])
+			// Plus as we push the value, the stack pointer must be incremented.
+			require.Equal(t, uint64(1), eng.currentStackPointer)
+		})
+	}
+}
+
+func TestAmd64Builder_handleGlobalSet(t *testing.T) {
+	const valueToSet uint64 = 12345
+	for i, tp := range []wasm.ValueType{
+		wasm.ValueTypeF32, wasm.ValueTypeF64, wasm.ValueTypeI32, wasm.ValueTypeI64,
+	} {
+		tp := tp
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+			// Setup the globals.
+			builder := requireNewBuilder(t)
+			globals := []*wasm.GlobalInstance{nil, {Val: 40, Type: &wasm.GlobalType{ValType: tp}}, nil}
+			builder.f = &wasm.FunctionInstance{ModuleInstance: &wasm.ModuleInstance{Globals: globals}}
+			_ = builder.locationStack.pushValueOnStack() // where we place the set target value below.
+			// Now emit the code.
+			builder.initializeReservedRegisters()
+			op := &wazeroir.OperationGlobalSet{Index: 1}
+			err := builder.handleGlobalSet(op)
+			require.NoError(t, err)
+			builder.returnFunction()
+
+			// Assemble.
+			code, err := builder.assemble()
+			require.NoError(t, err)
+			// Run code.
+			eng := newEngine()
+			eng.currentGlobalSliceAddress = uintptr(unsafe.Pointer(&globals[0]))
+			eng.push(valueToSet)
+			mem := newMemoryInst()
+			jitcall(
+				uintptr(unsafe.Pointer(&code[0])),
+				uintptr(unsafe.Pointer(eng)),
+				uintptr(unsafe.Pointer(&mem.Buffer[0])),
+			)
+			// The global value should be set to valueToSet.
+			require.Equal(t, valueToSet, globals[op.Index].Val)
+			// Plus we consumed the top of the stack, the stack pointer must be decremented.
+			require.Equal(t, uint64(0), eng.currentStackPointer)
+		})
+	}
+}