Repository: tetratelabs/wazero
Pull request #244 (2022-02-15 04:20:06)

jit(amd64): less instructions for memory bounds check: part 2
Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>


Patch:
diff --git a/wasm/jit/jit_amd64.go b/wasm/jit/jit_amd64.go
index 9efb37fe..034285c2 100644
--- a/wasm/jit/jit_amd64.go
+++ b/wasm/jit/jit_amd64.go
@@ -3907,30 +3907,30 @@ func (c *amd64Compiler) compileGe(o *wazeroir.OperationGe) error {
 // compileLoad implements compiler.compileLoad for the amd64 architecture.
 func (c *amd64Compiler) compileLoad(o *wazeroir.OperationLoad) error {
 	var (
-		isIntType        bool
-		movInst          obj.As
-		targetSizeInByte int64
+		isIntType         bool
+		movInst           obj.As
+		targetSizeInBytes int64
 	)
 	switch o.Type {
 	case wazeroir.UnsignedTypeI32:
 		isIntType = true
 		movInst = x86.AMOVL
-		targetSizeInByte = 32 / 8
+		targetSizeInBytes = 32 / 8
 	case wazeroir.UnsignedTypeI64:
 		isIntType = true
 		movInst = x86.AMOVQ
-		targetSizeInByte = 64 / 8
+		targetSizeInBytes = 64 / 8
 	case wazeroir.UnsignedTypeF32:
 		isIntType = false
 		movInst = x86.AMOVL
-		targetSizeInByte = 32 / 8
+		targetSizeInBytes = 32 / 8
 	case wazeroir.UnsignedTypeF64:
 		isIntType = false
 		movInst = x86.AMOVQ
-		targetSizeInByte = 64 / 8
+		targetSizeInBytes = 64 / 8
 	}
 
-	reg, err := c.setupMemoryOffset(o.Arg.Offset, targetSizeInByte)
+	reg, err := c.setupMemoryAccessCeil(o.Arg.Offset, targetSizeInBytes)
 	if err != nil {
 		return err
 	}
@@ -3944,6 +3944,8 @@ func (c *amd64Compiler) compileLoad(o *wazeroir.OperationLoad) error {
 		moveFromMemory.To.Reg = reg
 		moveFromMemory.From.Type = obj.TYPE_MEM
 		moveFromMemory.From.Reg = reservedRegisterForMemory
+		// because this is accessed as memory.Buffer[ceil-targetSizeInBytes: ceil]
+		moveFromMemory.From.Offset = -targetSizeInBytes
 		moveFromMemory.From.Index = reg
 		moveFromMemory.From.Scale = 1
 		c.addInstruction(moveFromMemory)
@@ -3960,6 +3962,8 @@ func (c *amd64Compiler) compileLoad(o *wazeroir.OperationLoad) error {
 		moveFromMemory.To.Type = obj.TYPE_REG
 		moveFromMemory.To.Reg = floatReg
 		moveFromMemory.From.Type = obj.TYPE_MEM
+		// because this is accessed as memory.Buffer[ceil-targetSizeInBytes: ceil]
+		moveFromMemory.From.Offset = -targetSizeInBytes
 		moveFromMemory.From.Reg = reservedRegisterForMemory
 		moveFromMemory.From.Index = reg
 		moveFromMemory.From.Scale = 1
@@ -3974,7 +3978,8 @@ func (c *amd64Compiler) compileLoad(o *wazeroir.OperationLoad) error {
 
 // compileLoad8 implements compiler.compileLoad8 for the amd64 architecture.
 func (c *amd64Compiler) compileLoad8(o *wazeroir.OperationLoad8) error {
-	reg, err := c.setupMemoryOffset(o.Arg.Offset, 1)
+	const targetSizeInBytes = 1
+	reg, err := c.setupMemoryAccessCeil(o.Arg.Offset, targetSizeInBytes)
 	if err != nil {
 		return err
 	}
@@ -3996,6 +4001,8 @@ func (c *amd64Compiler) compileLoad8(o *wazeroir.OperationLoad8) error {
 	moveFromMemory.To.Reg = reg
 	moveFromMemory.From.Type = obj.TYPE_MEM
 	moveFromMemory.From.Reg = reservedRegisterForMemory
+	// because this is accessed as memory.Buffer[ceil-targetSizeInBytes: ceil]
+	moveFromMemory.From.Offset = -targetSizeInBytes
 	moveFromMemory.From.Index = reg
 	moveFromMemory.From.Scale = 1
 	c.addInstruction(moveFromMemory)
@@ -4008,7 +4015,8 @@ func (c *amd64Compiler) compileLoad8(o *wazeroir.OperationLoad8) error {
 
 // compileLoad16 implements compiler.compileLoad16 for the amd64 architecture.
 func (c *amd64Compiler) compileLoad16(o *wazeroir.OperationLoad16) error {
-	reg, err := c.setupMemoryOffset(o.Arg.Offset, 16/8)
+	const targetSizeInBytes = 16 / 8
+	reg, err := c.setupMemoryAccessCeil(o.Arg.Offset, targetSizeInBytes)
 	if err != nil {
 		return err
 	}
@@ -4030,6 +4038,8 @@ func (c *amd64Compiler) compileLoad16(o *wazeroir.OperationLoad16) error {
 	moveFromMemory.To.Reg = reg
 	moveFromMemory.From.Type = obj.TYPE_MEM
 	moveFromMemory.From.Reg = reservedRegisterForMemory
+	// because this is accessed as memory.Buffer[ceil-targetSizeInBytes: ceil]
+	moveFromMemory.From.Offset = -targetSizeInBytes
 	moveFromMemory.From.Index = reg
 	moveFromMemory.From.Scale = 1
 	c.addInstruction(moveFromMemory)
@@ -4042,7 +4052,8 @@ func (c *amd64Compiler) compileLoad16(o *wazeroir.OperationLoad16) error {
 
 // compileLoad32 implements compiler.compileLoad32 for the amd64 architecture.
 func (c *amd64Compiler) compileLoad32(o *wazeroir.OperationLoad32) error {
-	reg, err := c.setupMemoryOffset(o.Arg.Offset, 32/8)
+	const targetSizeInBytes = 32 / 8
+	reg, err := c.setupMemoryAccessCeil(o.Arg.Offset, targetSizeInBytes)
 	if err != nil {
 		return err
 	}
@@ -4058,6 +4069,8 @@ func (c *amd64Compiler) compileLoad32(o *wazeroir.OperationLoad32) error {
 	moveFromMemory.To.Reg = reg
 	moveFromMemory.From.Type = obj.TYPE_MEM
 	moveFromMemory.From.Reg = reservedRegisterForMemory
+	// because this is accessed as memory.Buffer[ceil-targetSizeInBytes: ceil]
+	moveFromMemory.From.Offset = -targetSizeInBytes
 	moveFromMemory.From.Index = reg
 	moveFromMemory.From.Scale = 1
 	c.addInstruction(moveFromMemory)
@@ -4068,15 +4081,13 @@ func (c *amd64Compiler) compileLoad32(o *wazeroir.OperationLoad32) error {
 	return nil
 }
 
-// setupMemoryOffset pops the top value from the stack (called "base"), and returns the result of addition with
-// base and offsetArg, which we call "offset". The returned offsetRegister is the register number with the offset calculation value.
-// targetSizeInBytes is the original memory operation's target size in byte. For example, 4 = 32 / 8 for Load32 operation.
-// This is used for all Store* and Load* instructions.
+// setupMemoryAccessCeil pops the top value from the stack (called "base"), stores "base + offsetArg + targetSizeInBytes"
+// into a register, and returns the stored register. We call the result "ceil" because we access the memory
+// as memory.Buffer[ceil-targetSizeInBytes: ceil].
 //
-// Note that this also emits the instructions to check the out of bounds memory access. That means
-// if the base+offsetArg+targetSizeInBytes exceeds the memory size, we exit this function with
-// jitCallStatusCodeMemoryOutOfBounds status code since we read memory as [base+offsetArg: base+offsetArg+targetSizeInBytes].
-func (c *amd64Compiler) setupMemoryOffset(offsetArg uint32, targetSizeInBytes int64) (int16, error) {
+// Note: this also emits the instructions to check the out of bounds memory access.
+// In other words, if the ceil exceeds the memory size, the code exits with jitCallStatusCodeMemoryOutOfBounds status.
+func (c *amd64Compiler) setupMemoryAccessCeil(offsetArg uint32, targetSizeInBytes int64) (int16, error) {
 	base := c.locationStack.pop()
 	if err := c.ensureOnGeneralPurposeRegister(base); err != nil {
 		return 0, err
@@ -4118,16 +4129,6 @@ func (c *amd64Compiler) setupMemoryOffset(offsetArg uint32, targetSizeInBytes in
 
 	c.addSetJmpOrigins(okJmp)
 
-	// Finally we have to subtract targetSizeInBytes from the result register
-	// because we access memory with [offset: offset+targetSizeInBytes]
-	sub := c.newProg()
-	sub.As = x86.ASUBL
-	sub.To.Type = obj.TYPE_REG
-	sub.To.Reg = result
-	sub.From.Type = obj.TYPE_CONST
-	sub.From.Offset = targetSizeInBytes
-	c.addInstruction(sub)
-
 	c.locationStack.markRegisterUnused(result)
 	return result, nil
 }
@@ -4168,7 +4169,7 @@ func (c *amd64Compiler) moveToMemory(offsetConst uint32, moveInstruction obj.As,
 		return err
 	}
 
-	reg, err := c.setupMemoryOffset(offsetConst, targetSizeInByte)
+	reg, err := c.setupMemoryAccessCeil(offsetConst, targetSizeInByte)
 	if err != nil {
 		return nil
 	}
@@ -4179,6 +4180,8 @@ func (c *amd64Compiler) moveToMemory(offsetConst uint32, moveInstruction obj.As,
 	moveToMemory.From.Reg = val.register
 	moveToMemory.To.Type = obj.TYPE_MEM
 	moveToMemory.To.Reg = reservedRegisterForMemory
+	// because this is accessed as memory.Buffer[ceil-targetSizeInBytes: ceil]
+	moveToMemory.To.Offset = -targetSizeInByte
 	moveToMemory.To.Index = reg
 	moveToMemory.To.Scale = 1
 	c.addInstruction(moveToMemory)


Test patch:
diff --git a/wasm/jit/jit_amd64_test.go b/wasm/jit/jit_amd64_test.go
index bb321822..67a056c1 100644
--- a/wasm/jit/jit_amd64_test.go
+++ b/wasm/jit/jit_amd64_test.go
@@ -4774,7 +4774,7 @@ func TestAmd64Compiler_compile_min_max_copysign(t *testing.T) {
 	}
 }
 
-func TestAmd64Compiler_setupMemoryOffset(t *testing.T) {
+func TestAmd64Compiler_setupMemoryAccessCeil(t *testing.T) {
 	bases := []uint32{0, 1 << 5, 1 << 9, 1 << 10, 1 << 15, math.MaxUint32 - 1, math.MaxUint32}
 	offsets := []uint32{0,
 		1 << 10, 1 << 31, math.MaxInt32 - 1, math.MaxInt32 - 2, math.MaxInt32 - 3, math.MaxInt32 - 4,
@@ -4798,7 +4798,7 @@ func TestAmd64Compiler_setupMemoryOffset(t *testing.T) {
 					err = compiler.compileConstI32(&wazeroir.OperationConstI32{Value: base})
 					require.NoError(t, err)
 
-					reg, err := compiler.setupMemoryOffset(offset, targetSizeInByte)
+					reg, err := compiler.setupMemoryAccessCeil(offset, targetSizeInByte)
 					require.NoError(t, err)
 
 					compiler.locationStack.pushValueLocationOnRegister(reg)
@@ -4813,15 +4813,14 @@ func TestAmd64Compiler_setupMemoryOffset(t *testing.T) {
 					env.exec(code)
 
 					mem := env.memory()
-					baseOffset := int64(base) + int64(offset)
-					if ceil := baseOffset + int64(targetSizeInByte); int64(len(mem)) < ceil {
+					if ceil := int64(base) + int64(offset) + int64(targetSizeInByte); int64(len(mem)) < ceil {
 						// If the targe memory region's ceil exceeds the length of memory, we must exit the function
 						// with jitCallStatusCodeMemoryOutOfBounds status code.
 						require.Equal(t, jitCallStatusCodeMemoryOutOfBounds, env.jitStatus())
 					} else {
 						require.Equal(t, jitCallStatusCodeReturned, env.jitStatus())
 						require.Equal(t, uint64(1), env.stackPointer())
-						require.Equal(t, uint64(baseOffset), env.stackTopAsUint64())
+						require.Equal(t, uint64(ceil), env.stackTopAsUint64())
 					}
 				})
 			}