Repository: tetratelabs/wazero
Pull request #93 (2021-12-24 05:32:15)

jit: support for store instructions

mathetake: ![img](https://media.giphy.com/media/11EjiLDatd0syA/giphy.gif)

Patch:
diff --git a/wasm/jit/jit_amd64.go b/wasm/jit/jit_amd64.go
index 3d6e27b7..3d7d5764 100644
--- a/wasm/jit/jit_amd64.go
+++ b/wasm/jit/jit_amd64.go
@@ -114,13 +114,21 @@ func (e *engine) compileWasmFunction(f *wasm.FunctionInstance) (*compiledWasmFun
 				return nil, fmt.Errorf("error handling load16 operation: %w", err)
 			}
 		case *wazeroir.OperationStore:
-			return nil, fmt.Errorf("unsupported operation in JIT compiler: %v", o)
+			if err := builder.handleStore(o); err != nil {
+				return nil, fmt.Errorf("error handling store operation: %w", err)
+			}
 		case *wazeroir.OperationStore8:
-			return nil, fmt.Errorf("unsupported operation in JIT compiler: %v", o)
+			if err := builder.handleStore8(o); err != nil {
+				return nil, fmt.Errorf("error handling store8 operation: %w", err)
+			}
 		case *wazeroir.OperationStore16:
-			return nil, fmt.Errorf("unsupported operation in JIT compiler: %v", o)
+			if err := builder.handleStore16(o); err != nil {
+				return nil, fmt.Errorf("error handling store16 operation: %w", err)
+			}
 		case *wazeroir.OperationStore32:
-			return nil, fmt.Errorf("unsupported operation in JIT compiler: %v", o)
+			if err := builder.handleStore32(o); err != nil {
+				return nil, fmt.Errorf("error handling store32 operation: %w", err)
+			}
 		case *wazeroir.OperationMemorySize:
 			builder.handleMemorySize()
 		case *wazeroir.OperationMemoryGrow:
@@ -1289,6 +1297,68 @@ func (b *amd64Builder) handleLoad32(o *wazeroir.OperationLoad32) error {
 	return nil
 }
 
+func (b *amd64Builder) handleStore(o *wazeroir.OperationStore) error {
+	var movInst obj.As
+	switch o.Type {
+	case wazeroir.UnsignedTypeI32, wazeroir.UnsignedTypeF32:
+		movInst = x86.AMOVL
+	case wazeroir.UnsignedTypeI64, wazeroir.UnsignedTypeF64:
+		movInst = x86.AMOVQ
+	}
+	return b.moveToMemory(o.Arg.Offest, movInst)
+}
+
+func (b *amd64Builder) handleStore8(o *wazeroir.OperationStore8) error {
+	return b.moveToMemory(o.Arg.Offest, x86.AMOVB)
+}
+
+func (b *amd64Builder) handleStore16(o *wazeroir.OperationStore16) error {
+	return b.moveToMemory(o.Arg.Offest, x86.AMOVW)
+}
+
+func (b *amd64Builder) handleStore32(o *wazeroir.OperationStore32) error {
+	return b.moveToMemory(o.Arg.Offest, x86.AMOVL)
+}
+
+func (b *amd64Builder) moveToMemory(offsetConst uint32, moveInstruction obj.As) error {
+	val := b.locationStack.pop()
+	if err := b.ensureOnGeneralPurposeRegister(val); err != nil {
+		return err
+	}
+
+	base := b.locationStack.pop()
+	if err := b.ensureOnGeneralPurposeRegister(base); err != nil {
+		return err
+	}
+
+	// Then we have to calculate the offset on the memory region.
+	addOffsetToBase := b.newProg()
+	addOffsetToBase.As = x86.AADDL // 32-bit!
+	addOffsetToBase.To.Type = obj.TYPE_REG
+	addOffsetToBase.To.Reg = base.register
+	addOffsetToBase.From.Type = obj.TYPE_CONST
+	addOffsetToBase.From.Offset = int64(offsetConst)
+	b.addInstruction(addOffsetToBase)
+
+	// TODO: Emit instructions here to check memory out of bounds as
+	// potentially it would be an security risk.
+
+	moveToMemory := b.newProg()
+	moveToMemory.As = moveInstruction
+	moveToMemory.From.Type = obj.TYPE_REG
+	moveToMemory.From.Reg = val.register
+	moveToMemory.To.Type = obj.TYPE_MEM
+	moveToMemory.To.Reg = reservedRegisterForMemory
+	moveToMemory.To.Index = base.register
+	moveToMemory.To.Scale = 1
+	b.addInstruction(moveToMemory)
+
+	// We no longer need both the value and base registers.
+	b.locationStack.releaseRegister(val)
+	b.locationStack.releaseRegister(base)
+	return nil
+}
+
 func (b *amd64Builder) handleMemoryGrow() {
 	b.callBuiltinFunctionFromConstIndex(builtinFunctionIndexMemoryGrow)
 }


Test patch:
diff --git a/wasm/jit/jit_amd64_test.go b/wasm/jit/jit_amd64_test.go
index e66a958e..ed8f9bec 100644
--- a/wasm/jit/jit_amd64_test.go
+++ b/wasm/jit/jit_amd64_test.go
@@ -1684,6 +1684,213 @@ func TestAmd64Builder_handleLoad32(t *testing.T) {
 	require.Equal(t, original+1, uint32(eng.stack[eng.stackPointer-1]))
 }
 
+func TestAmd64Builder_handleStore(t *testing.T) {
+	for i, tp := range []wazeroir.UnsignedType{
+		wazeroir.UnsignedTypeI32,
+		wazeroir.UnsignedTypeI64,
+		wazeroir.UnsignedTypeF32,
+		wazeroir.UnsignedTypeF64,
+	} {
+		tp := tp
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+			eng := newEngine()
+			builder := requireNewBuilder(t)
+			builder.initializeReservedRegisters()
+
+			// Before store operations, we must push the base offset, and the store target values.
+			const baseOffset = 100 // For testing. Arbitrary number is fine.
+			base := builder.locationStack.pushValueOnStack()
+			eng.stack[base.stackPointer] = baseOffset
+			storeTargetValue := uint64(math.MaxUint64)
+			storeTarget := builder.locationStack.pushValueOnStack()
+			eng.stack[storeTarget.stackPointer] = storeTargetValue
+			switch tp {
+			case wazeroir.UnsignedTypeI32, wazeroir.UnsignedTypeF32:
+				storeTarget.setRegisterType(generalPurposeRegisterTypeInt)
+			case wazeroir.UnsignedTypeI64, wazeroir.UnsignedTypeF64:
+				storeTarget.setRegisterType(generalPurposeRegisterTypeFloat)
+			}
+
+			// Emit the memory load instructions.
+			o := &wazeroir.OperationStore{Type: tp, Arg: &wazeroir.MemoryImmediate{Offest: 361}}
+			err := builder.handleStore(o)
+			require.NoError(t, err)
+
+			// At this point, two values are popped so the stack pointer must be zero.
+			require.Equal(t, uint64(0), builder.locationStack.sp)
+			// Plus there should be no used registers.
+			require.Len(t, builder.locationStack.usedRegisters, 0)
+
+			// Compile.
+			builder.returnFunction()
+			code, err := builder.compile()
+			require.NoError(t, err)
+
+			// Run code.
+			mem := newMemoryInst()
+			jitcall(
+				uintptr(unsafe.Pointer(&code[0])),
+				uintptr(unsafe.Pointer(eng)),
+				uintptr(unsafe.Pointer(&mem.Buffer[0])),
+			)
+
+			// All the values are popped, so the stack pointer must be zero.
+			require.Equal(t, uint64(0), eng.stackPointer)
+			// Check the stored value.
+			offset := o.Arg.Offest + baseOffset
+			switch o.Type {
+			case wazeroir.UnsignedTypeI32, wazeroir.UnsignedTypeF32:
+				v := binary.LittleEndian.Uint32(mem.Buffer[offset : offset+4])
+				require.Equal(t, uint32(storeTargetValue), v)
+				// The trailing bytes must be intact since this is 32-bit mov.
+				v = binary.LittleEndian.Uint32(mem.Buffer[offset+4 : offset+8])
+				require.Equal(t, uint32(0), v)
+			case wazeroir.UnsignedTypeI64, wazeroir.UnsignedTypeF64:
+				v := binary.LittleEndian.Uint64(mem.Buffer[offset : offset+8])
+				require.Equal(t, storeTargetValue, v)
+			}
+		})
+	}
+}
+
+func TestAmd64Builder_handleStore8(t *testing.T) {
+	eng := newEngine()
+	builder := requireNewBuilder(t)
+	builder.initializeReservedRegisters()
+
+	// Before store operations, we must push the base offset, and the store target values.
+	const baseOffset = 100 // For testing. Arbitrary number is fine.
+	base := builder.locationStack.pushValueOnStack()
+	eng.stack[base.stackPointer] = baseOffset
+	storeTargetValue := uint64(0x12_34_56_78_9a_bc_ef_01) // For testing. Arbitrary number is fine.
+	storeTarget := builder.locationStack.pushValueOnStack()
+	eng.stack[storeTarget.stackPointer] = storeTargetValue
+	storeTarget.setRegisterType(generalPurposeRegisterTypeInt)
+
+	// Emit the memory load instructions.
+	o := &wazeroir.OperationStore8{Arg: &wazeroir.MemoryImmediate{Offest: 361}}
+	err := builder.handleStore8(o)
+	require.NoError(t, err)
+
+	// At this point, two values are popped so the stack pointer must be zero.
+	require.Equal(t, uint64(0), builder.locationStack.sp)
+	// Plus there should be no used registers.
+	require.Len(t, builder.locationStack.usedRegisters, 0)
+
+	// Compile.
+	builder.returnFunction()
+	code, err := builder.compile()
+	require.NoError(t, err)
+
+	// Run code.
+	mem := newMemoryInst()
+	jitcall(
+		uintptr(unsafe.Pointer(&code[0])),
+		uintptr(unsafe.Pointer(eng)),
+		uintptr(unsafe.Pointer(&mem.Buffer[0])),
+	)
+
+	// All the values are popped, so the stack pointer must be zero.
+	require.Equal(t, uint64(0), eng.stackPointer)
+	// Check the stored value.
+	offset := o.Arg.Offest + baseOffset
+	require.Equal(t, byte(storeTargetValue), mem.Buffer[offset])
+	// The trailing bytes must be intact since this is only moving one byte.
+	require.Equal(t, []byte{0, 0, 0, 0, 0, 0, 0}, mem.Buffer[offset+1:offset+8])
+}
+
+func TestAmd64Builder_handleStore16(t *testing.T) {
+	eng := newEngine()
+	builder := requireNewBuilder(t)
+	builder.initializeReservedRegisters()
+
+	// Before store operations, we must push the base offset, and the store target values.
+	const baseOffset = 100 // For testing. Arbitrary number is fine.
+	base := builder.locationStack.pushValueOnStack()
+	eng.stack[base.stackPointer] = baseOffset
+	storeTargetValue := uint64(0x12_34_56_78_9a_bc_ef_01) // For testing. Arbitrary number is fine.
+	storeTarget := builder.locationStack.pushValueOnStack()
+	eng.stack[storeTarget.stackPointer] = storeTargetValue
+	storeTarget.setRegisterType(generalPurposeRegisterTypeInt)
+
+	// Emit the memory load instructions.
+	o := &wazeroir.OperationStore16{Arg: &wazeroir.MemoryImmediate{Offest: 361}}
+	err := builder.handleStore16(o)
+	require.NoError(t, err)
+
+	// At this point, two values are popped so the stack pointer must be zero.
+	require.Equal(t, uint64(0), builder.locationStack.sp)
+	// Plus there should be no used registers.
+	require.Len(t, builder.locationStack.usedRegisters, 0)
+
+	// Compile.
+	builder.returnFunction()
+	code, err := builder.compile()
+	require.NoError(t, err)
+
+	// Run code.
+	mem := newMemoryInst()
+	jitcall(
+		uintptr(unsafe.Pointer(&code[0])),
+		uintptr(unsafe.Pointer(eng)),
+		uintptr(unsafe.Pointer(&mem.Buffer[0])),
+	)
+
+	// All the values are popped, so the stack pointer must be zero.
+	require.Equal(t, uint64(0), eng.stackPointer)
+	// Check the stored value.
+	offset := o.Arg.Offest + baseOffset
+	require.Equal(t, uint16(storeTargetValue), binary.LittleEndian.Uint16(mem.Buffer[offset:]))
+	// The trailing bytes must be intact since this is only moving 2 byte.
+	require.Equal(t, []byte{0, 0, 0, 0, 0, 0}, mem.Buffer[offset+2:offset+8])
+}
+
+func TestAmd64Builder_handleStore32(t *testing.T) {
+	eng := newEngine()
+	builder := requireNewBuilder(t)
+	builder.initializeReservedRegisters()
+
+	// Before store operations, we must push the base offset, and the store target values.
+	const baseOffset = 100 // For testing. Arbitrary number is fine.
+	base := builder.locationStack.pushValueOnStack()
+	eng.stack[base.stackPointer] = baseOffset
+	storeTargetValue := uint64(0x12_34_56_78_9a_bc_ef_01) // For testing. Arbitrary number is fine.
+	storeTarget := builder.locationStack.pushValueOnStack()
+	eng.stack[storeTarget.stackPointer] = storeTargetValue
+	storeTarget.setRegisterType(generalPurposeRegisterTypeInt)
+
+	// Emit the memory load instructions.
+	o := &wazeroir.OperationStore32{Arg: &wazeroir.MemoryImmediate{Offest: 361}}
+	err := builder.handleStore32(o)
+	require.NoError(t, err)
+
+	// At this point, two values are popped so the stack pointer must be zero.
+	require.Equal(t, uint64(0), builder.locationStack.sp)
+	// Plus there should be no used registers.
+	require.Len(t, builder.locationStack.usedRegisters, 0)
+
+	// Compile.
+	builder.returnFunction()
+	code, err := builder.compile()
+	require.NoError(t, err)
+
+	// Run code.
+	mem := newMemoryInst()
+	jitcall(
+		uintptr(unsafe.Pointer(&code[0])),
+		uintptr(unsafe.Pointer(eng)),
+		uintptr(unsafe.Pointer(&mem.Buffer[0])),
+	)
+
+	// All the values are popped, so the stack pointer must be zero.
+	require.Equal(t, uint64(0), eng.stackPointer)
+	// Check the stored value.
+	offset := o.Arg.Offest + baseOffset
+	require.Equal(t, uint32(storeTargetValue), binary.LittleEndian.Uint32(mem.Buffer[offset:]))
+	// The trailing bytes must be intact since this is only moving 4 byte.
+	require.Equal(t, []byte{0, 0, 0, 0}, mem.Buffer[offset+4:offset+8])
+}
+
 func TestAmd64Builder_handleMemoryGrow(t *testing.T) {
 	builder := requireNewBuilder(t)