Repository: tetratelabs/wazero
Pull request #86 (2021-12-23 06:45:20)

jit: support memory.load instruction


Patch:
diff --git a/wasm/jit/jit_amd64.go b/wasm/jit/jit_amd64.go
index 69f83bde..e2ba42ee 100644
--- a/wasm/jit/jit_amd64.go
+++ b/wasm/jit/jit_amd64.go
@@ -98,7 +98,9 @@ func (e *engine) compileWasmFunction(f *wasm.FunctionInstance) (*compiledWasmFun
 				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)
+			if err := builder.handleLoad(o); err != nil {
+				return nil, fmt.Errorf("error handling load operation: %w", err)
+			}
 		case *wazeroir.OperationLoad8:
 			return nil, fmt.Errorf("unsupported operation in JIT compiler: %v", o)
 		case *wazeroir.OperationLoad16:
@@ -328,7 +330,7 @@ func (b *amd64Builder) handleSwap(o *wazeroir.OperationSwap) error {
 	// If x1 is on the conditional register, we must move it to a gp
 	// register before swap.
 	if x1.onConditionalRegister() {
-		if err := b.moveConditionalToGPRegister(x1); err != nil {
+		if err := b.moveConditionalToGeneralPurposeRegister(x1); err != nil {
 			return err
 		}
 	}
@@ -473,7 +475,7 @@ func (b *amd64Builder) handleGlobalSet(o *wazeroir.OperationGlobalSet) error {
 			return err
 		}
 	} else if val.onConditionalRegister() {
-		if err := b.moveConditionalToGPRegister(val); err != nil {
+		if err := b.moveConditionalToGeneralPurposeRegister(val); err != nil {
 			return err
 		}
 	}
@@ -768,7 +770,7 @@ func (b *amd64Builder) emitDropRange(r *wazeroir.InclusiveRange) error {
 			if topIsConditional {
 				// If the top is conditional, and it's not target of drop,
 				// we must assign it to the register before we emit any instructions here.
-				if err := b.moveConditionalToGPRegister(top); err != nil {
+				if err := b.moveConditionalToGeneralPurposeRegister(top); err != nil {
 					return err
 				}
 				topIsConditional = false
@@ -801,7 +803,7 @@ func (b *amd64Builder) handleSelect() error {
 
 	// Ensure the conditional value lives in a gp register.
 	if c.onConditionalRegister() {
-		if err := b.moveConditionalToGPRegister(c); err != nil {
+		if err := b.moveConditionalToGeneralPurposeRegister(c); err != nil {
 			return err
 		}
 	} else if c.onStack() {
@@ -932,7 +934,7 @@ func (b *amd64Builder) handleAdd(o *wazeroir.OperationAdd) error {
 			return err
 		}
 	} else if x2.onConditionalRegister() {
-		if err := b.moveConditionalToGPRegister(x2); err != nil {
+		if err := b.moveConditionalToGeneralPurposeRegister(x2); err != nil {
 			return err
 		}
 	}
@@ -993,7 +995,7 @@ func (b *amd64Builder) handleSub(o *wazeroir.OperationSub) error {
 			return err
 		}
 	} else if x2.onConditionalRegister() {
-		if err := b.moveConditionalToGPRegister(x2); err != nil {
+		if err := b.moveConditionalToGeneralPurposeRegister(x2); err != nil {
 			return err
 		}
 	}
@@ -1059,7 +1061,7 @@ func (b *amd64Builder) handleLe(o *wazeroir.OperationLe) error {
 			return err
 		}
 	} else if x2.onConditionalRegister() {
-		if err := b.moveConditionalToGPRegister(x2); err != nil {
+		if err := b.moveConditionalToGeneralPurposeRegister(x2); err != nil {
 			return err
 		}
 	}
@@ -1096,6 +1098,90 @@ func (b *amd64Builder) handleLe(o *wazeroir.OperationLe) error {
 	return nil
 }
 
+func (b *amd64Builder) handleLoad(o *wazeroir.OperationLoad) error {
+	base := b.locationStack.pop()
+	if base.onStack() {
+		if err := b.moveStackToRegisterWithAllocation(generalPurposeRegisterTypeInt, base); err != nil {
+			return err
+		}
+	} else if base.onConditionalRegister() {
+		if err := b.moveConditionalToGeneralPurposeRegister(base); err != nil {
+			return err
+		}
+	}
+
+	// At this point, base's value is on the integer general purpose reg.
+	// We reuse the register below, so we alias it here for readability.
+	reg := base.register
+
+	// 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 = reg
+	addOffsetToBase.From.Type = obj.TYPE_CONST
+	addOffsetToBase.From.Offset = int64(o.Arg.Offest)
+	b.addInstruction(addOffsetToBase)
+
+	// TODO: Emit instructions here to check memory out of bounds as
+	// potentially it would be an security risk.
+
+	var (
+		isIntType bool
+		movInst   obj.As
+	)
+	switch o.Type {
+	case wazeroir.SignLessTypeI32:
+		isIntType = true
+		movInst = x86.AMOVL
+	case wazeroir.SignLessTypeI64:
+		isIntType = true
+		movInst = x86.AMOVQ
+	case wazeroir.SignLessTypeF32:
+		isIntType = false
+		movInst = x86.AMOVL
+	case wazeroir.SignLessTypeF64:
+		isIntType = false
+		movInst = x86.AMOVQ
+	}
+
+	if isIntType {
+		// For integer types, read the corresponding bytes from the offset to the memory
+		// and store the value to the int register.
+		moveFromMemory := b.newProg()
+		moveFromMemory.As = movInst
+		moveFromMemory.To.Type = obj.TYPE_REG
+		moveFromMemory.To.Reg = reg
+		moveFromMemory.From.Type = obj.TYPE_MEM
+		moveFromMemory.From.Reg = reservedRegisterForMemory
+		moveFromMemory.From.Index = reg
+		moveFromMemory.From.Scale = 1
+		b.addInstruction(moveFromMemory)
+		top := b.locationStack.pushValueOnRegister(reg)
+		top.setValueType(o.Type)
+	} else {
+		// For float types, we read the value to the float register.
+		floatReg, err := b.allocateRegister(generalPurposeRegisterTypeFloat)
+		if err != nil {
+			return err
+		}
+		moveFromMemory := b.newProg()
+		moveFromMemory.As = movInst
+		moveFromMemory.To.Type = obj.TYPE_REG
+		moveFromMemory.To.Reg = floatReg
+		moveFromMemory.From.Type = obj.TYPE_MEM
+		moveFromMemory.From.Reg = reservedRegisterForMemory
+		moveFromMemory.From.Index = reg
+		moveFromMemory.From.Scale = 1
+		b.addInstruction(moveFromMemory)
+		top := b.locationStack.pushValueOnRegister(floatReg)
+		top.setValueType(o.Type)
+		// We no longer need the int register so mark it unused.
+		b.locationStack.markRegisterUnused(reg)
+	}
+	return nil
+}
+
 func (b *amd64Builder) handleMemoryGrow() {
 	b.callBuiltinFunctionFromConstIndex(builtinFunctionIndexMemoryGrow)
 }
@@ -1250,11 +1336,11 @@ func (b *amd64Builder) moveStackToRegister(loc *valueLocation) {
 	b.addInstruction(prog)
 }
 
-func (b *amd64Builder) moveConditionalToGPRegister(loc *valueLocation) error {
+func (b *amd64Builder) moveConditionalToGeneralPurposeRegister(loc *valueLocation) error {
 	// Get the free register.
 	reg, ok := b.locationStack.takeFreeRegister(generalPurposeRegisterTypeInt)
 	if !ok {
-		// This in theory should never be reached as moveConditionalToGPRegister
+		// This in theory should never be reached as moveConditionalToGeneralPurposeRegister
 		// is called right after comparison operations, meaning that
 		// at least two registers are free at the moment.
 		return fmt.Errorf("conditional register mov requires a free register")
diff --git a/wasm/jit/jit_value_location_amd64.go b/wasm/jit/jit_value_location_amd64.go
index f48d252c..e04e4aa8 100644
--- a/wasm/jit/jit_value_location_amd64.go
+++ b/wasm/jit/jit_value_location_amd64.go
@@ -202,11 +202,15 @@ func (s *valueLocationStack) peek() (loc *valueLocation) {
 }
 
 func (s *valueLocationStack) releaseRegister(loc *valueLocation) {
-	delete(s.usedRegisters, loc.register)
+	s.markRegisterUnused(loc.register)
 	loc.register = -1
 	loc.conditionalRegister = conditionalRegisterStateUnset
 }
 
+func (s *valueLocationStack) markRegisterUnused(reg int16) {
+	delete(s.usedRegisters, reg)
+}
+
 func (s *valueLocationStack) markRegisterUsed(reg int16) {
 	s.usedRegisters[reg] = struct{}{}
 }


Test patch:
diff --git a/wasm/jit/jit_amd64_test.go b/wasm/jit/jit_amd64_test.go
index 52d470fd..56742735 100644
--- a/wasm/jit/jit_amd64_test.go
+++ b/wasm/jit/jit_amd64_test.go
@@ -1092,7 +1092,7 @@ func TestAmd64Builder_handleLe(t *testing.T) {
 			// to the stack.
 			top := builder.locationStack.peek()
 			require.True(t, top.onConditionalRegister() && !top.onRegister())
-			err = builder.moveConditionalToGPRegister(top)
+			err = builder.moveConditionalToGeneralPurposeRegister(top)
 			require.NoError(t, err)
 			require.True(t, !top.onConditionalRegister() && top.onRegister())
 			builder.releaseRegisterToStack(top)
@@ -1150,7 +1150,7 @@ func TestAmd64Builder_handleLe(t *testing.T) {
 			// to the stack.
 			top := builder.locationStack.peek()
 			require.True(t, top.onConditionalRegister() && !top.onRegister())
-			err = builder.moveConditionalToGPRegister(top)
+			err = builder.moveConditionalToGeneralPurposeRegister(top)
 			require.NoError(t, err)
 			require.True(t, !top.onConditionalRegister() && top.onRegister())
 			builder.releaseRegisterToStack(top)
@@ -1402,6 +1402,103 @@ func TestAmd64Builder_handleCall(t *testing.T) {
 	})
 }
 
+func TestAmd64Builder_handleLoad(t *testing.T) {
+	for i, tp := range []wazeroir.SignLessType{
+		wazeroir.SignLessTypeI32,
+		wazeroir.SignLessTypeI64,
+		wazeroir.SignLessTypeF32,
+		wazeroir.SignLessTypeF64,
+	} {
+		tp := tp
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+			eng := newEngine()
+			builder := requireNewBuilder(t)
+			builder.initializeReservedRegisters()
+			// Before load operations, we must push the base offset value.
+			const baseOffset = 100 // For testing. Arbitrary number is fine.
+			base := builder.locationStack.pushValueOnStack()
+			eng.stack[base.stackPointer] = baseOffset
+
+			// Emit the memory load instructions.
+			o := &wazeroir.OperationLoad{Type: tp, Arg: &wazeroir.MemoryImmediate{Offest: 361}}
+			err := builder.handleLoad(o)
+			require.NoError(t, err)
+
+			// At this point, the loaded value must be on top of the stack, and placed on a register.
+			loadedValue := builder.locationStack.peek()
+			require.Equal(t, o.Type, loadedValue.valueType)
+			require.True(t, loadedValue.onRegister())
+
+			// Double the loaded value in order to verify the behavior.
+			var addInst obj.As
+			switch tp {
+			case wazeroir.SignLessTypeI32:
+				require.True(t, isIntRegister(loadedValue.register))
+				addInst = x86.AADDL
+			case wazeroir.SignLessTypeI64:
+				require.True(t, isIntRegister(loadedValue.register))
+				addInst = x86.AADDQ
+			case wazeroir.SignLessTypeF32:
+				require.True(t, isFloatRegister(loadedValue.register))
+				addInst = x86.AADDSS
+			case wazeroir.SignLessTypeF64:
+				require.True(t, isFloatRegister(loadedValue.register))
+				addInst = x86.AADDSD
+			}
+			doubleLoadedValue := builder.newProg()
+			doubleLoadedValue.As = addInst
+			doubleLoadedValue.To.Type = obj.TYPE_REG
+			doubleLoadedValue.To.Reg = loadedValue.register
+			doubleLoadedValue.From.Type = obj.TYPE_REG
+			doubleLoadedValue.From.Reg = loadedValue.register
+			builder.addInstruction(doubleLoadedValue)
+
+			// We need to write the result back to the memory stack.
+			builder.releaseRegisterToStack(loadedValue)
+
+			// Compile.
+			builder.returnFunction()
+			code, err := builder.compile()
+			require.NoError(t, err)
+
+			// Place the load target value to the memory.
+			mem := newMemoryInst()
+			targetRegion := mem.Buffer[baseOffset+o.Arg.Offest:]
+			var expValue uint64
+			switch tp {
+			case wazeroir.SignLessTypeI32:
+				original := uint32(100)
+				binary.LittleEndian.PutUint32(targetRegion, original)
+				expValue = uint64(original * 2)
+			case wazeroir.SignLessTypeI64:
+				original := uint64(math.MaxUint32 + 123) // The value exceeds 32-bit.
+				binary.LittleEndian.PutUint64(targetRegion, original)
+				expValue = original * 2
+			case wazeroir.SignLessTypeF32:
+				original := float32(1.234)
+				binary.LittleEndian.PutUint32(targetRegion, math.Float32bits(original))
+				expValue = uint64(math.Float32bits(original * 2))
+			case wazeroir.SignLessTypeF64:
+				original := float64(math.MaxFloat32 + 100.1) // The value exceeds 32-bit.
+				binary.LittleEndian.PutUint64(targetRegion, math.Float64bits(original))
+				expValue = math.Float64bits(original * 2)
+			}
+
+			// Run code.
+			jitcall(
+				uintptr(unsafe.Pointer(&code[0])),
+				uintptr(unsafe.Pointer(eng)),
+				uintptr(unsafe.Pointer(&mem.Buffer[0])),
+			)
+
+			// Load instruction must push the loaded value to the top of the stack,
+			// so the stack pointer must be incremented.
+			require.Equal(t, uint64(1), eng.currentStackPointer)
+			require.Equal(t, expValue, eng.stack[eng.currentStackPointer-1])
+		})
+	}
+}
+
 func TestAmd64Builder_handleMemoryGrow(t *testing.T) {
 	builder := requireNewBuilder(t)