Repository: tetratelabs/wazero
Pull request #96 (2021-12-30 02:53:39)

jit: support for less than equal for floats.
Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>


Patch:
diff --git a/wasm/jit/jit_amd64.go b/wasm/jit/jit_amd64.go
index d8e3bd05..63bc54a5 100644
--- a/wasm/jit/jit_amd64.go
+++ b/wasm/jit/jit_amd64.go
@@ -1038,65 +1038,57 @@ func (b *amd64Builder) handleSub(o *wazeroir.OperationSub) error {
 }
 
 func (b *amd64Builder) handleLe(o *wazeroir.OperationLe) error {
+	x2 := b.locationStack.pop()
+	if err := b.ensureOnGeneralPurposeRegister(x2); err != nil {
+		return err
+	}
+
+	x1 := b.locationStack.pop()
+	if err := b.ensureOnGeneralPurposeRegister(x1); err != nil {
+		return err
+	}
+
+	// Emit the compare instruction.
+	prog := b.newProg()
+	prog.From.Type = obj.TYPE_REG
+	prog.To.Type = obj.TYPE_REG
 	var resultConditionState conditionalRegisterState
-	var instruction obj.As
-	var tp generalPurposeRegisterType
 	switch o.Type {
 	case wazeroir.SignedTypeInt32:
 		resultConditionState = conditionalRegisterStateLE
-		instruction = x86.ACMPL
-		tp = generalPurposeRegisterTypeInt
+		prog.As = x86.ACMPL
+		prog.From.Reg = x1.register
+		prog.To.Reg = x2.register
 	case wazeroir.SignedTypeUint32:
 		resultConditionState = conditionalRegisterStateBE
-		instruction = x86.ACMPL
-		tp = generalPurposeRegisterTypeInt
+		prog.As = x86.ACMPL
+		prog.From.Reg = x1.register
+		prog.To.Reg = x2.register
 	case wazeroir.SignedTypeInt64:
 		resultConditionState = conditionalRegisterStateLE
-		instruction = x86.ACMPQ
-		tp = generalPurposeRegisterTypeInt
+		prog.As = x86.ACMPQ
+		prog.From.Reg = x1.register
+		prog.To.Reg = x2.register
 	case wazeroir.SignedTypeUint64:
 		resultConditionState = conditionalRegisterStateBE
-		instruction = x86.ACMPQ
-		tp = generalPurposeRegisterTypeInt
+		prog.As = x86.ACMPQ
+		prog.From.Reg = x1.register
+		prog.To.Reg = x2.register
 	case wazeroir.SignedTypeFloat32:
-		tp = generalPurposeRegisterTypeFloat
-		panic("add test!")
+		resultConditionState = conditionalRegisterStateBE
+		prog.As = x86.ACOMISS
+		prog.From.Reg = x2.register
+		prog.To.Reg = x1.register
 	case wazeroir.SignedTypeFloat64:
-		tp = generalPurposeRegisterTypeFloat
-		panic("add test!")
-	}
-
-	x2 := b.locationStack.pop()
-	if x2.onStack() {
-		if err := b.moveStackToRegisterWithAllocation(tp, x2); err != nil {
-			return err
-		}
-	} else if x2.onConditionalRegister() {
-		if err := b.moveConditionalToGeneralPurposeRegister(x2); err != nil {
-			return err
-		}
-	}
-
-	x1 := b.locationStack.pop()
-	if x1.onStack() {
-		if err := b.moveStackToRegisterWithAllocation(tp, x1); err != nil {
-			return err
-		}
-	} else if x1.onConditionalRegister() {
-		// This shouldn't happen as the conditional
-		// register must be on top of the stack.
-		panic("a bug in jit compiler")
+		resultConditionState = conditionalRegisterStateBE
+		prog.As = x86.ACOMISD
+		prog.From.Reg = x2.register
+		prog.To.Reg = x1.register
 	}
-
-	// Compare: set the flag based on x1-x2.
-	prog := b.newProg()
-	prog.As = instruction
-	prog.From.Type = obj.TYPE_REG
-	prog.From.Reg = x1.register
-	prog.To.Type = obj.TYPE_REG
-	prog.To.Reg = x2.register
 	b.addInstruction(prog)
 
+	// TODO: emit NaN value handings for floats.
+
 	// We no longer need x1,x2 register after cmp operation here,
 	// so we release it.
 	b.locationStack.releaseRegister(x1)


Test patch:
diff --git a/wasm/jit/jit_amd64_test.go b/wasm/jit/jit_amd64_test.go
index 31cf3c91..70877ceb 100644
--- a/wasm/jit/jit_amd64_test.go
+++ b/wasm/jit/jit_amd64_test.go
@@ -1127,6 +1127,10 @@ func TestAmd64Builder_handleLe(t *testing.T) {
 			{x1: -1, x2: 100, signed: false, exp: false},
 			{x1: 100, x2: 200, signed: true, exp: true},
 			{x1: 200, x2: 100, signed: true, exp: false},
+			{x1: 1 << 56, x2: 100, signed: true, exp: false},
+			{x1: 1 << 56, x2: 1 << 61, signed: true, exp: true},
+			{x1: math.MaxInt64, x2: 100, signed: true, exp: false},
+			{x1: math.MinInt64, x2: 100, signed: true, exp: true},
 		} {
 			var o *wazeroir.OperationLe
 			if tc.signed {
@@ -1176,6 +1180,140 @@ func TestAmd64Builder_handleLe(t *testing.T) {
 			}
 		}
 	})
+	t.Run("float32", func(t *testing.T) {
+		for i, tc := range []struct {
+			x1, x2 float32
+			exp    bool
+		}{
+			{x1: 100, x2: -1.1, exp: false},
+			{x1: -1, x2: 100, exp: true},
+			{x1: 100, x2: 200, exp: true},
+			{x1: 100.01234124, x2: 100.01234124, exp: true},
+			{x1: 100.01234124, x2: -100.01234124, exp: false},
+			{x1: 200.12315, x2: 100, exp: false},
+			{x1: float32(math.Inf(1)), x2: 100, exp: false},
+			{x1: float32(math.Inf(-1)), x2: 100, exp: true},
+		} {
+			t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+				// Prepare operands.
+				builder := requireNewBuilder(t)
+				builder.initializeReservedRegisters()
+				err := builder.handleConstF32(&wazeroir.OperationConstF32{Value: tc.x1})
+				require.NoError(t, err)
+				x1 := builder.locationStack.peek()
+				err = builder.handleConstF32(&wazeroir.OperationConstF32{Value: tc.x2})
+				require.NoError(t, err)
+				x2 := builder.locationStack.peek()
+
+				// Emit the Le instructions,
+				err = builder.handleLe(&wazeroir.OperationLe{Type: wazeroir.SignedTypeFloat32})
+				require.NoError(t, err)
+
+				// At this point, these registers must be consumed.
+				require.NotContains(t, builder.locationStack.usedRegisters, x1.register)
+				require.NotContains(t, builder.locationStack.usedRegisters, x2.register)
+				// Plus the result must be pushed.
+				require.Equal(t, uint64(1), builder.locationStack.sp)
+
+				// To verify the behavior, we push the flag value
+				// to the stack.
+				flag := builder.locationStack.peek()
+				require.True(t, flag.onConditionalRegister() && !flag.onRegister())
+				err = builder.moveConditionalToGeneralPurposeRegister(flag)
+				require.NoError(t, err)
+				require.True(t, !flag.onConditionalRegister() && flag.onRegister())
+				builder.releaseRegisterToStack(flag)
+				builder.returnFunction()
+
+				// Assemble.
+				code, err := builder.compile()
+				require.NoError(t, err)
+				// Run code.
+				eng := newEngine()
+				mem := newMemoryInst()
+				jitcall(
+					uintptr(unsafe.Pointer(&code[0])),
+					uintptr(unsafe.Pointer(eng)),
+					uintptr(unsafe.Pointer(&mem.Buffer[0])),
+				)
+				// Check the stack.
+				require.Equal(t, uint64(1), eng.stackPointer)
+				if tc.exp {
+					require.Equal(t, uint64(1), eng.stack[eng.stackPointer-1])
+				} else {
+					require.Equal(t, uint64(0), eng.stack[eng.stackPointer-1])
+				}
+			})
+		}
+	})
+	t.Run("float64", func(t *testing.T) {
+		for i, tc := range []struct {
+			x1, x2 float64
+			exp    bool
+		}{
+			{x1: 100, x2: -1.1, exp: false},
+			{x1: -1, x2: 100, exp: true},
+			{x1: 100, x2: 200, exp: true},
+			{x1: 100.01234124, x2: 100.01234124, exp: true},
+			{x1: 100.01234124, x2: -100.01234124, exp: false},
+			{x1: 200.12315, x2: 100, exp: false},
+			{x1: 6.8719476736e+10 /* = 1 << 36 */, x2: 100, exp: false},
+			{x1: 6.8719476736e+10 /* = 1 << 36 */, x2: 1.37438953472e+11 /* = 1 << 37*/, exp: true},
+			{x1: math.Inf(1), x2: 100, exp: false},
+			{x1: math.Inf(-1), x2: 100, exp: true},
+		} {
+			t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+				// Prepare operands.
+				builder := requireNewBuilder(t)
+				builder.initializeReservedRegisters()
+				err := builder.handleConstF64(&wazeroir.OperationConstF64{Value: tc.x1})
+				require.NoError(t, err)
+				x1 := builder.locationStack.peek()
+				err = builder.handleConstF64(&wazeroir.OperationConstF64{Value: tc.x2})
+				require.NoError(t, err)
+				x2 := builder.locationStack.peek()
+
+				// Emit the Le instructions,
+				err = builder.handleLe(&wazeroir.OperationLe{Type: wazeroir.SignedTypeFloat64})
+				require.NoError(t, err)
+
+				// At this point, these registers must be consumed.
+				require.NotContains(t, builder.locationStack.usedRegisters, x1.register)
+				require.NotContains(t, builder.locationStack.usedRegisters, x2.register)
+				// Plus the result must be pushed.
+				require.Equal(t, uint64(1), builder.locationStack.sp)
+
+				// To verify the behavior, we push the flag value
+				// to the stack.
+				flag := builder.locationStack.peek()
+				require.True(t, flag.onConditionalRegister() && !flag.onRegister())
+				err = builder.moveConditionalToGeneralPurposeRegister(flag)
+				require.NoError(t, err)
+				require.True(t, !flag.onConditionalRegister() && flag.onRegister())
+				builder.releaseRegisterToStack(flag)
+				builder.returnFunction()
+
+				// Assemble.
+				code, err := builder.compile()
+				require.NoError(t, err)
+				// Run code.
+				eng := newEngine()
+				mem := newMemoryInst()
+				jitcall(
+					uintptr(unsafe.Pointer(&code[0])),
+					uintptr(unsafe.Pointer(eng)),
+					uintptr(unsafe.Pointer(&mem.Buffer[0])),
+				)
+				// Check the stack.
+				require.Equal(t, uint64(1), eng.stackPointer)
+				if tc.exp {
+					require.Equal(t, uint64(1), eng.stack[eng.stackPointer-1])
+				} else {
+					require.Equal(t, uint64(0), eng.stack[eng.stackPointer-1])
+				}
+			})
+		}
+	})
 }
 
 func TestAmd64Builder_handleSub(t *testing.T) {