Repository: tetratelabs/wazero
Pull request #126 (2022-01-10 10:49:52)

jit: completes numeric instructions.

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

Patch:
diff --git a/wasm/jit/compiler.go b/wasm/jit/compiler.go
index 8b3f7189..1a0bb082 100644
--- a/wasm/jit/compiler.go
+++ b/wasm/jit/compiler.go
@@ -53,6 +53,11 @@ type compiler interface {
 	compileFConvertFromI(o *wazeroir.OperationFConvertFromI) error
 	compileF32DemoteFromF64() error
 	compileF64PromoteFromF32() error
+	compileI32ReinterpretFromF32() error
+	compileI64ReinterpretFromF64() error
+	compileF32ReinterpretFromI32() error
+	compileF64ReinterpretFromI64() error
+	compileExtend(o *wazeroir.OperationExtend) error
 	compileEq(o *wazeroir.OperationEq) error
 	compileNe(o *wazeroir.OperationNe) error
 	compileEqz(o *wazeroir.OperationEqz) error
diff --git a/wasm/jit/engine.go b/wasm/jit/engine.go
index d3e2daa2..d230f4e2 100644
--- a/wasm/jit/engine.go
+++ b/wasm/jit/engine.go
@@ -595,13 +595,16 @@ func (e *engine) compileWasmFunction(f *wasm.FunctionInstance) (*compiledWasmFun
 			err = compiler.compileF32DemoteFromF64()
 		case *wazeroir.OperationF64PromoteFromF32:
 			err = compiler.compileF64PromoteFromF32()
-		case *wazeroir.OperationI32ReinterpretFromF32,
-			*wazeroir.OperationI64ReinterpretFromF64,
-			*wazeroir.OperationF32ReinterpretFromI32,
-			*wazeroir.OperationF64ReinterpretFromI64:
-			err = fmt.Errorf("unsupported operation")
+		case *wazeroir.OperationI32ReinterpretFromF32:
+			err = compiler.compileI32ReinterpretFromF32()
+		case *wazeroir.OperationI64ReinterpretFromF64:
+			err = compiler.compileI64ReinterpretFromF64()
+		case *wazeroir.OperationF32ReinterpretFromI32:
+			err = compiler.compileI64ReinterpretFromF64()
+		case *wazeroir.OperationF64ReinterpretFromI64:
+			err = compiler.compileF64ReinterpretFromI64()
 		case *wazeroir.OperationExtend:
-			err = fmt.Errorf("unsupported operation")
+			err = compiler.compileExtend(o)
 		}
 		if err != nil {
 			return nil, fmt.Errorf("failed to compile operation %s: %w", op.Kind().String(), err)
diff --git a/wasm/jit/jit_amd64.go b/wasm/jit/jit_amd64.go
index 3b6d682f..3bfdd42a 100644
--- a/wasm/jit/jit_amd64.go
+++ b/wasm/jit/jit_amd64.go
@@ -2497,13 +2497,13 @@ func (c *amd64Compiler) emitSignedI64TruncFromFloat(isFloat32Bit bool) error {
 // the corresponding float value. This is equivalent to float32(uint32(x)), float32(int32(x)), etc in Go.
 func (c *amd64Compiler) compileFConvertFromI(o *wazeroir.OperationFConvertFromI) (err error) {
 	if o.OutputType == wazeroir.Float32 && o.InputType == wazeroir.SignedInt32 {
-		err = c.emitSimpleIntToFloatConversion(x86.ACVTSL2SS) // = CVTSI2SS for 32bit int
+		err = c.emitSimpleConversion(x86.ACVTSL2SS, generalPurposeRegisterTypeFloat) // = CVTSI2SS for 32bit int
 	} else if o.OutputType == wazeroir.Float32 && o.InputType == wazeroir.SignedInt64 {
-		err = c.emitSimpleIntToFloatConversion(x86.ACVTSQ2SS) // = CVTSI2SS for 64bit int
+		err = c.emitSimpleConversion(x86.ACVTSQ2SS, generalPurposeRegisterTypeFloat) // = CVTSI2SS for 64bit int
 	} else if o.OutputType == wazeroir.Float64 && o.InputType == wazeroir.SignedInt32 {
-		err = c.emitSimpleIntToFloatConversion(x86.ACVTSL2SD) // = CVTSI2SD for 32bit int
+		err = c.emitSimpleConversion(x86.ACVTSL2SD, generalPurposeRegisterTypeFloat) // = CVTSI2SD for 32bit int
 	} else if o.OutputType == wazeroir.Float64 && o.InputType == wazeroir.SignedInt64 {
-		err = c.emitSimpleIntToFloatConversion(x86.ACVTSQ2SD) // = CVTSI2SD for 64bit int
+		err = c.emitSimpleConversion(x86.ACVTSQ2SD, generalPurposeRegisterTypeFloat) // = CVTSI2SD for 64bit int
 	} else if o.OutputType == wazeroir.Float32 && o.InputType == wazeroir.SignedUint32 {
 		// See the following link for why we use 64bit conversion for unsigned 32bit integer sources:
 		// https://stackoverflow.com/questions/41495498/fpu-operations-generated-by-gcc-during-casting-integer-to-float.
@@ -2513,10 +2513,10 @@ func (c *amd64Compiler) compileFConvertFromI(o *wazeroir.OperationFConvertFromI)
 		// >> not an unsigned integer like you have here. So what gives? Well, a 64-bit processor has 64-bit wide
 		// >> registers available, so the unsigned 32-bit input values can be stored as signed 64-bit intermediate values,
 		// >> which allows CVTSI2SS to be used after all.
-		err = c.emitSimpleIntToFloatConversion(x86.ACVTSQ2SS) // = CVTSI2SS for 64bit int.
+		err = c.emitSimpleConversion(x86.ACVTSQ2SS, generalPurposeRegisterTypeFloat) // = CVTSI2SS for 64bit int.
 	} else if o.OutputType == wazeroir.Float64 && o.InputType == wazeroir.SignedUint32 {
 		// For the same reason above, we use 64bit conversion for unsigned 32bit.
-		err = c.emitSimpleIntToFloatConversion(x86.ACVTSQ2SD) // = CVTSI2SD for 64bit int.
+		err = c.emitSimpleConversion(x86.ACVTSQ2SD, generalPurposeRegisterTypeFloat) // = CVTSI2SD for 64bit int.
 	} else if o.OutputType == wazeroir.Float32 && o.InputType == wazeroir.SignedUint64 {
 		err = c.emitUnsignedInt64ToFloatConversion(true)
 	} else if o.OutputType == wazeroir.Float64 && o.InputType == wazeroir.SignedUint64 {
@@ -2687,15 +2687,16 @@ func (c *amd64Compiler) emitUnsignedInt64ToFloatConversion(isFloat32bit bool) er
 	return nil
 }
 
-// emitSimpleIntToFloatConversion pops a flaot type from the stack, and applies the
-// given instruction on it, and push the integer result onto the stack.
-func (c *amd64Compiler) emitSimpleIntToFloatConversion(convInstruction obj.As) error {
+// emitSimpleConversion pops a value type from the stack, and applies the
+// given instruction on it, and push the result onto a register of the given type.
+func (c *amd64Compiler) emitSimpleConversion(convInstruction obj.As, destinationRegisterType generalPurposeRegisterType) error {
+	// TODO: it is not necessary to place the origin on a register (i.e. ok to directly move the memory allocated value).
 	origin := c.locationStack.pop()
 	if err := c.ensureOnGeneralPurposeRegister(origin); err != nil {
 		return err
 	}
 
-	dest, err := c.allocateRegister(generalPurposeRegisterTypeFloat)
+	dest, err := c.allocateRegister(destinationRegisterType)
 	if err != nil {
 		return err
 	}
@@ -2710,7 +2711,7 @@ func (c *amd64Compiler) emitSimpleIntToFloatConversion(convInstruction obj.As) e
 
 	c.locationStack.markRegisterUnused(origin.register)
 	loc := c.locationStack.pushValueOnRegister(dest)
-	loc.setRegisterType(generalPurposeRegisterTypeFloat)
+	loc.setRegisterType(destinationRegisterType)
 	return nil
 }
 
@@ -2750,6 +2751,70 @@ func (c *amd64Compiler) compileF64PromoteFromF32() error {
 	return nil
 }
 
+// compileI32ReinterpretFromF32 adds instructions to reinterpret the 32-bit float on top of the stack
+// as a 32-bit integer by preserving the bit representation. If the value is on the stack,
+// this is no-op as there is nothing to do for converting type.
+func (c *amd64Compiler) compileI32ReinterpretFromF32() error {
+	if c.locationStack.peek().onStack() {
+		return nil
+	}
+	return c.emitSimpleConversion(x86.AMOVL, generalPurposeRegisterTypeInt)
+}
+
+// compileI64ReinterpretFromF64 adds instructions to reinterpret the 64-bit float on top of the stack
+// as a 64-bit integer by preserving the bit representation. If the value is on the stack,
+// this is no-op as there is nothing to do for converting type.
+func (c *amd64Compiler) compileI64ReinterpretFromF64() error {
+	if c.locationStack.peek().onStack() {
+		return nil
+	}
+	return c.emitSimpleConversion(x86.AMOVQ, generalPurposeRegisterTypeInt)
+}
+
+// compileF32ReinterpretFromI32 adds instructions to reinterpret the 32-bit int on top of the stack
+// as a 32-bit float by preserving the bit representation. If the value is on the stack,
+// this is no-op as there is nothing to do for converting type.
+func (c *amd64Compiler) compileF32ReinterpretFromI32() error {
+	if c.locationStack.peek().onStack() {
+		return nil
+	}
+	return c.emitSimpleConversion(x86.AMOVL, generalPurposeRegisterTypeFloat)
+}
+
+// compileF32ReinterpretFromI32 adds instructions to reinterpret the 64-bit int on top of the stack
+// as a 64-bit float by preserving the bit representation. If the value is on the stack,
+// this is no-op as there is nothing to do for converting type.
+func (c *amd64Compiler) compileF64ReinterpretFromI64() error {
+	if c.locationStack.peek().onStack() {
+		return nil
+	}
+	return c.emitSimpleConversion(x86.AMOVQ, generalPurposeRegisterTypeFloat)
+}
+
+// compileExtend adds instructions to extend the 32-bit signed or unsigned int on top of the stack
+// as a 64-bit integer of coressponding signedness. For unsigned case, this is just reinterpreting the
+// underlying bit pattern as 64-bit integer. For signed case, this is sign-extension which preserves the
+// original integer's sign.
+func (c *amd64Compiler) compileExtend(o *wazeroir.OperationExtend) error {
+	target := c.locationStack.peek() // Note this is peek!
+	if err := c.ensureOnGeneralPurposeRegister(target); err != nil {
+		return err
+	}
+
+	extend := c.newProg()
+	if o.Signed {
+		extend.As = x86.AMOVLQSX // = MOVSXD https://www.felixcloutier.com/x86/movsx:movsxd
+	} else {
+		extend.As = x86.AMOVQ
+	}
+	extend.From.Type = obj.TYPE_REG
+	extend.From.Reg = target.register
+	extend.To.Type = obj.TYPE_REG
+	extend.To.Reg = target.register
+	c.addInstruction(extend)
+	return nil
+}
+
 func (c *amd64Compiler) compileEq(o *wazeroir.OperationEq) error {
 	return c.emitEqOrNe(o.Type, true)
 }


Test patch:
diff --git a/wasm/jit/jit_amd64_test.go b/wasm/jit/jit_amd64_test.go
index e43bc0a2..a01a9a8e 100644
--- a/wasm/jit/jit_amd64_test.go
+++ b/wasm/jit/jit_amd64_test.go
@@ -4000,6 +4000,139 @@ func TestAmd64Compiler_compileF64PromoteFromF32(t *testing.T) {
 	}
 }
 
+func TestAmd64Compiler_compileReinterpret(t *testing.T) {
+	for _, kind := range []wazeroir.OperationKind{
+		wazeroir.OperationKindF32ReinterpretFromI32,
+		wazeroir.OperationKindF64ReinterpretFromI64,
+		wazeroir.OperationKindI32ReinterpretFromF32,
+		wazeroir.OperationKindI64ReinterpretFromF64,
+	} {
+		kind := kind
+		t.Run(kind.String(), func(t *testing.T) {
+			for _, originOnStack := range []bool{false, true} {
+				originOnStack := originOnStack
+				t.Run(fmt.Sprintf("%v", originOnStack), func(t *testing.T) {
+					for _, v := range []uint64{
+						0, 1, 1 << 16, 1 << 31, 1 << 32, 1 << 63,
+						math.MaxInt32, math.MaxUint32, math.MaxUint64,
+					} {
+						v := v
+						t.Run(fmt.Sprintf("%d", v), func(t *testing.T) {
+							compiler := requireNewCompiler(t)
+							compiler.initializeReservedRegisters()
+
+							eng := newEngine()
+							if originOnStack {
+								loc := compiler.locationStack.pushValueOnStack()
+								eng.stack[loc.stackPointer] = v
+								eng.stackPointer = 1
+							}
+
+							var is32Bit bool
+							var err error
+							switch kind {
+							case wazeroir.OperationKindF32ReinterpretFromI32:
+								is32Bit = true
+								if !originOnStack {
+									err = compiler.compileConstI32(&wazeroir.OperationConstI32{Value: uint32(v)})
+									require.NoError(t, err)
+								}
+								err = compiler.compileF32ReinterpretFromI32()
+								require.NoError(t, err)
+							case wazeroir.OperationKindF64ReinterpretFromI64:
+								if !originOnStack {
+									err = compiler.compileConstI64(&wazeroir.OperationConstI64{Value: v})
+									require.NoError(t, err)
+								}
+								err = compiler.compileF64ReinterpretFromI64()
+								require.NoError(t, err)
+							case wazeroir.OperationKindI32ReinterpretFromF32:
+								is32Bit = true
+								if !originOnStack {
+									err = compiler.compileConstF32(&wazeroir.OperationConstF32{Value: math.Float32frombits(uint32(v))})
+									require.NoError(t, err)
+								}
+								err = compiler.compileI32ReinterpretFromF32()
+								require.NoError(t, err)
+							case wazeroir.OperationKindI64ReinterpretFromF64:
+								if !originOnStack {
+									err = compiler.compileConstF64(&wazeroir.OperationConstF64{Value: math.Float64frombits(v)})
+									require.NoError(t, err)
+								}
+								err = compiler.compileI64ReinterpretFromF64()
+								require.NoError(t, err)
+							default:
+								t.Fail()
+							}
+
+							// To verify the behavior, we release the value
+							// to the stack.
+							compiler.releaseAllRegistersToStack()
+							compiler.returnFunction()
+
+							// Generate and run the code under test.
+							code, _, err := compiler.generate()
+							require.NoError(t, err)
+							jitcall(uintptr(unsafe.Pointer(&code[0])), uintptr(unsafe.Pointer(eng)), 0)
+
+							// Reinterpret must preserve the bit-pattern.
+							if is32Bit {
+								require.Equal(t, uint32(v), stackTopAsUint32(eng))
+							} else {
+								require.Equal(t, v, stackTopAsUint64(eng))
+							}
+						})
+					}
+				})
+			}
+		})
+	}
+}
+
+func TestAmd64Compiler_compileExtend(t *testing.T) {
+	for _, signed := range []bool{false, true} {
+		signed := signed
+		t.Run(fmt.Sprintf("signed=%v", signed), func(t *testing.T) {
+			for _, v := range []uint32{
+				0, 1, 1 << 14, 1 << 31, math.MaxUint32, 0xFFFFFFFF, math.MaxInt32,
+			} {
+				v := v
+				t.Run(fmt.Sprintf("%v", v), func(t *testing.T) {
+					compiler := requireNewCompiler(t)
+					compiler.initializeReservedRegisters()
+
+					// Setup the promote target.
+					err := compiler.compileConstI32(&wazeroir.OperationConstI32{Value: v})
+					require.NoError(t, err)
+
+					err = compiler.compileExtend(&wazeroir.OperationExtend{Signed: signed})
+					require.NoError(t, err)
+
+					// To verify the behavior, we release the value
+					// to the stack.
+					compiler.releaseAllRegistersToStack()
+					compiler.returnFunction()
+
+					// Generate and run the code under test.
+					code, _, err := compiler.generate()
+					require.NoError(t, err)
+					eng := newEngine()
+					jitcall(uintptr(unsafe.Pointer(&code[0])), uintptr(unsafe.Pointer(eng)), 0)
+
+					require.Equal(t, uint64(1), eng.stackPointer)
+					if signed {
+						expected := int64(int32(v))
+						require.Equal(t, expected, stackTopAsInt64(eng))
+					} else {
+						expected := uint64(uint32(v))
+						require.Equal(t, expected, stackTopAsUint64(eng))
+					}
+				})
+			}
+		})
+	}
+}
+
 func TestAmd64Compiler_compileITruncFromF(t *testing.T) {
 	for _, tc := range []struct {
 		outputType wazeroir.SignedInt