Repository: tetratelabs/wazero
Pull request #118 (2022-01-06 07:09:10)

jit: support for and, or, xor, shl, shr, rotl and rotr operations
Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>

mathetake: Implemented Rotr and Rotl as well! https://github.com/tetratelabs/wazero/pull/118/commits/bbeaec5baa4578740593babfcb0b49fc63f5e658

Patch:
diff --git a/wasm/jit/compiler.go b/wasm/jit/compiler.go
index a974943f..d10d548f 100644
--- a/wasm/jit/compiler.go
+++ b/wasm/jit/compiler.go
@@ -31,6 +31,13 @@ type compiler interface {
 	compilePopcnt(o *wazeroir.OperationPopcnt) error
 	compileDiv(o *wazeroir.OperationDiv) error
 	compileRem(o *wazeroir.OperationRem) error
+	compileAnd(o *wazeroir.OperationAnd) error
+	compileOr(o *wazeroir.OperationOr) error
+	compileXor(o *wazeroir.OperationXor) error
+	compileShl(o *wazeroir.OperationShl) error
+	compileShr(o *wazeroir.OperationShr) error
+	compileRotl(o *wazeroir.OperationRotl) error
+	compileRotr(o *wazeroir.OperationRotr) 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 f803e35f..7920c16d 100644
--- a/wasm/jit/engine.go
+++ b/wasm/jit/engine.go
@@ -549,19 +549,19 @@ func (e *engine) compileWasmFunction(f *wasm.FunctionInstance) (*compiledWasmFun
 		case *wazeroir.OperationRem:
 			err = compiler.compileRem(o)
 		case *wazeroir.OperationAnd:
-			err = fmt.Errorf("unsupported operation")
+			err = compiler.compileAnd(o)
 		case *wazeroir.OperationOr:
-			err = fmt.Errorf("unsupported operation")
+			err = compiler.compileOr(o)
 		case *wazeroir.OperationXor:
-			err = fmt.Errorf("unsupported operation")
+			err = compiler.compileXor(o)
 		case *wazeroir.OperationShl:
-			err = fmt.Errorf("unsupported operation")
+			err = compiler.compileShl(o)
 		case *wazeroir.OperationShr:
-			err = fmt.Errorf("unsupported operation")
+			err = compiler.compileShr(o)
 		case *wazeroir.OperationRotl:
-			err = fmt.Errorf("unsupported operation")
+			err = compiler.compileRotl(o)
 		case *wazeroir.OperationRotr:
-			err = fmt.Errorf("unsupported operation")
+			err = compiler.compileRotr(o)
 		case *wazeroir.OperationAbs:
 			err = fmt.Errorf("unsupported operation")
 		case *wazeroir.OperationNeg:
diff --git a/wasm/jit/jit_amd64.go b/wasm/jit/jit_amd64.go
index 7f87755c..92707921 100644
--- a/wasm/jit/jit_amd64.go
+++ b/wasm/jit/jit_amd64.go
@@ -1312,31 +1312,131 @@ func (c *amd64Compiler) performDivisionOnInts(is32Bit bool, signed bool) error {
 // on the top two values of float type on the stack, placing the result back onto the stack.
 // For example, stack [..., 1.0, 4.0] results in [..., 0.25].
 func (c *amd64Compiler) compileDivForFloats(is32Bit bool) error {
+	if is32Bit {
+		return c.emitSimpleBinaryOp(x86.ADIVSS)
+	} else {
+		return c.emitSimpleBinaryOp(x86.ADIVSD)
+	}
+}
+
+// compileAnd emits instructions to perform an "and" operation on
+// top two values on the stack, and pushes the result.
+func (c *amd64Compiler) compileAnd(o *wazeroir.OperationAnd) (err error) {
+	switch o.Type {
+	case wazeroir.UnsignedInt32:
+		err = c.emitSimpleBinaryOp(x86.AANDL)
+	case wazeroir.UnsignedInt64:
+		err = c.emitSimpleBinaryOp(x86.AANDQ)
+	}
+	return
+}
+
+// compileOr emits instructions to perform an "or" operation on
+// top two values on the stack, and pushes the result.
+func (c *amd64Compiler) compileOr(o *wazeroir.OperationOr) (err error) {
+	switch o.Type {
+	case wazeroir.UnsignedInt32:
+		err = c.emitSimpleBinaryOp(x86.AORL)
+	case wazeroir.UnsignedInt64:
+		err = c.emitSimpleBinaryOp(x86.AORQ)
+	}
+	return
+}
+
+// compileXor emits instructions to perform an xor operation on
+// top two values on the stack, and pushes the result.
+func (c *amd64Compiler) compileXor(o *wazeroir.OperationXor) (err error) {
+	switch o.Type {
+	case wazeroir.UnsignedInt32:
+		err = c.emitSimpleBinaryOp(x86.AXORL)
+	case wazeroir.UnsignedInt64:
+		err = c.emitSimpleBinaryOp(x86.AXORQ)
+	}
+	return
+}
+
+// compileShl emits instructions to perform a shift-left operation on
+// top two values on the stack, and pushes the result.
+func (c *amd64Compiler) compileShl(o *wazeroir.OperationShl) (err error) {
+	switch o.Type {
+	case wazeroir.UnsignedInt32:
+		err = c.emitSimpleBinaryOp(x86.ASHLL)
+	case wazeroir.UnsignedInt64:
+		err = c.emitSimpleBinaryOp(x86.ASHLQ)
+	}
+	return
+}
+
+// compileShr emits instructions to perform a shift-right operation on
+// top two values on the stack, and pushes the result.
+func (c *amd64Compiler) compileShr(o *wazeroir.OperationShr) (err error) {
+	switch o.Type {
+	case wazeroir.SignedInt32:
+		err = c.emitSimpleBinaryOp(x86.ASARL)
+	case wazeroir.SignedInt64:
+		err = c.emitSimpleBinaryOp(x86.ASARQ)
+	case wazeroir.SignedUint32:
+		err = c.emitSimpleBinaryOp(x86.ASHRL)
+	case wazeroir.SignedUint64:
+		err = c.emitSimpleBinaryOp(x86.ASHRQ)
+	}
+	return
+}
+
+// compileRotl emits instructions to perform a rotate-left operation on
+// top two values on the stack, and pushes the result.
+func (c *amd64Compiler) compileRotl(o *wazeroir.OperationRotl) (err error) {
+	switch o.Type {
+	case wazeroir.UnsignedInt32:
+		err = c.emitSimpleBinaryOp(x86.AROLL)
+	case wazeroir.UnsignedInt64:
+		err = c.emitSimpleBinaryOp(x86.AROLQ)
+	}
+	return
+}
+
+// compileRotr emits instructions to perform a rotate-right operation on
+// top two values on the stack, and pushes the result.
+func (c *amd64Compiler) compileRotr(o *wazeroir.OperationRotr) (err error) {
+	switch o.Type {
+	case wazeroir.UnsignedInt32:
+		err = c.emitSimpleBinaryOp(x86.ARORL)
+	case wazeroir.UnsignedInt64:
+		err = c.emitSimpleBinaryOp(x86.ARORQ)
+	}
+	return
+}
+
+// emitSimpleBinaryOp emits instructions to pop two values from the stack
+// and perform the given instruction on these two values and push the result
+// onto the stack.
+func (c *amd64Compiler) emitSimpleBinaryOp(instruction obj.As) error {
 	x2 := c.locationStack.pop()
 	if err := c.ensureOnGeneralPurposeRegister(x2); err != nil {
 		return err
 	}
 
-	x1 := c.locationStack.peek() // Note this is peek!
+	x1 := c.locationStack.pop()
 	if err := c.ensureOnGeneralPurposeRegister(x1); err != nil {
 		return err
 	}
 
-	div := c.newProg()
-	div.From.Type = obj.TYPE_REG
-	div.From.Reg = x2.register
-	div.To.Type = obj.TYPE_REG
-	div.To.Reg = x1.register
-	if is32Bit {
-		div.As = x86.ADIVSS
-	} else {
-		div.As = x86.ADIVSD
-	}
-	c.addInstruction(div)
+	inst := c.newProg()
+	inst.From.Type = obj.TYPE_REG
+	inst.From.Reg = x2.register
+	inst.To.Type = obj.TYPE_REG
+	inst.To.Reg = x1.register
+	inst.As = instruction
+	c.addInstruction(inst)
 
-	// We consumed x2 register after DIV operation here,
+	// We consumed x2 register after the operation here,
 	// so we release it.
 	c.locationStack.releaseRegister(x2)
+
+	// We already stored the result in the register used by x1
+	// so we record it.
+	result := c.locationStack.pushValueOnRegister(x1.register)
+	result.setRegisterType(x1.registerType())
 	return nil
 }
 


Test patch:
diff --git a/wasm/jit/jit_amd64_test.go b/wasm/jit/jit_amd64_test.go
index 1fa90501..89d86d75 100644
--- a/wasm/jit/jit_amd64_test.go
+++ b/wasm/jit/jit_amd64_test.go
@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"io"
 	"math"
+	"math/bits"
 	"reflect"
 	"runtime"
 	"sync"
@@ -2959,6 +2960,181 @@ func getDivisionByZeroErrorRecoverFunc(t *testing.T) func() {
 	}
 }
 
+func TestAmd64Compiler_compile_and_or_xor_shl_shr_rotl_rotr(t *testing.T) {
+	for _, tc := range []struct {
+		name string
+		op   wazeroir.Operation
+	}{
+		{name: "and-32-bit", op: &wazeroir.OperationAnd{Type: wazeroir.UnsignedInt32}},
+		{name: "and-64-bit", op: &wazeroir.OperationAnd{Type: wazeroir.UnsignedInt64}},
+		{name: "or-32-bit", op: &wazeroir.OperationOr{Type: wazeroir.UnsignedInt32}},
+		{name: "or-64-bit", op: &wazeroir.OperationOr{Type: wazeroir.UnsignedInt64}},
+		{name: "xor-32-bit", op: &wazeroir.OperationXor{Type: wazeroir.UnsignedInt32}},
+		{name: "xor-64-bit", op: &wazeroir.OperationXor{Type: wazeroir.UnsignedInt64}},
+		{name: "shl-32-bit", op: &wazeroir.OperationShl{Type: wazeroir.UnsignedInt32}},
+		{name: "shl-64-bit", op: &wazeroir.OperationShl{Type: wazeroir.UnsignedInt64}},
+		{name: "shr-signed-32-bit", op: &wazeroir.OperationShr{Type: wazeroir.SignedInt32}},
+		{name: "shr-signed-64-bit", op: &wazeroir.OperationShr{Type: wazeroir.SignedInt64}},
+		{name: "shr-unsigned-32-bit", op: &wazeroir.OperationShr{Type: wazeroir.SignedUint32}},
+		{name: "shr-unsigned-64-bit", op: &wazeroir.OperationShr{Type: wazeroir.SignedUint64}},
+		{name: "rotl-32-bit", op: &wazeroir.OperationRotl{Type: wazeroir.UnsignedInt32}},
+		{name: "rotl-64-bit", op: &wazeroir.OperationRotl{Type: wazeroir.UnsignedInt64}},
+		{name: "rotr-32-bit", op: &wazeroir.OperationRotr{Type: wazeroir.UnsignedInt32}},
+		{name: "rotr-64-bit", op: &wazeroir.OperationRotr{Type: wazeroir.UnsignedInt64}},
+	} {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			for i, vs := range []struct {
+				x1, x2 uint64
+			}{
+				{x1: 0, x2: 0},
+				{x1: 0, x2: 1},
+				{x1: 1, x2: 0},
+				{x1: 1, x2: 1},
+				{x1: 1 << 31, x2: 1},
+				{x1: 1, x2: 1 << 31},
+				{x1: 1 << 31, x2: 1 << 31},
+				{x1: 1 << 63, x2: 1},
+				{x1: 1, x2: 1 << 63},
+				{x1: 1 << 63, x2: 1 << 63},
+			} {
+				t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+					compiler := requireNewCompiler(t)
+					compiler.initializeReservedRegisters()
+
+					var is32Bit bool
+					var expectedValue uint64
+					var compileOperationFunc func()
+					switch o := tc.op.(type) {
+					case *wazeroir.OperationAnd:
+						compileOperationFunc = func() {
+							err := compiler.compileAnd(o)
+							require.NoError(t, err)
+						}
+						is32Bit = o.Type == wazeroir.UnsignedInt32
+						if is32Bit {
+							expectedValue = uint64(uint32(vs.x1) & uint32(vs.x2))
+						} else {
+							expectedValue = vs.x1 & vs.x2
+						}
+					case *wazeroir.OperationOr:
+						compileOperationFunc = func() {
+							err := compiler.compileOr(o)
+							require.NoError(t, err)
+						}
+						is32Bit = o.Type == wazeroir.UnsignedInt32
+						if is32Bit {
+							expectedValue = uint64(uint32(vs.x1) | uint32(vs.x2))
+						} else {
+							expectedValue = vs.x1 | vs.x2
+						}
+					case *wazeroir.OperationXor:
+						compileOperationFunc = func() {
+							err := compiler.compileXor(o)
+							require.NoError(t, err)
+						}
+						is32Bit = o.Type == wazeroir.UnsignedInt32
+						if is32Bit {
+							expectedValue = uint64(uint32(vs.x1) ^ uint32(vs.x2))
+						} else {
+							expectedValue = vs.x1 ^ vs.x2
+						}
+					case *wazeroir.OperationShl:
+						compileOperationFunc = func() {
+							err := compiler.compileShl(o)
+							require.NoError(t, err)
+						}
+						is32Bit = o.Type == wazeroir.UnsignedInt32
+						if is32Bit {
+							expectedValue = uint64(uint32(vs.x1) << uint32(vs.x2%32))
+						} else {
+							expectedValue = vs.x1 << (vs.x2 % 64)
+						}
+					case *wazeroir.OperationShr:
+						compileOperationFunc = func() {
+							err := compiler.compileShr(o)
+							require.NoError(t, err)
+						}
+						is32Bit = o.Type == wazeroir.SignedInt32 || o.Type == wazeroir.SignedUint32
+						switch o.Type {
+						case wazeroir.SignedInt32:
+							expectedValue = uint64(int32(vs.x1) >> (uint32(vs.x2) % 32))
+						case wazeroir.SignedInt64:
+							expectedValue = uint64(int64(vs.x1) >> (vs.x2 % 64))
+						case wazeroir.SignedUint32:
+							expectedValue = uint64(uint32(vs.x1) >> (uint32(vs.x2) % 32))
+						case wazeroir.SignedUint64:
+							expectedValue = vs.x1 >> (vs.x2 % 64)
+						}
+					case *wazeroir.OperationRotl:
+						compileOperationFunc = func() {
+							err := compiler.compileRotl(o)
+							require.NoError(t, err)
+						}
+						is32Bit = o.Type == wazeroir.UnsignedInt32
+						if is32Bit {
+							expectedValue = uint64(bits.RotateLeft32(uint32(vs.x1), int(vs.x2)))
+						} else {
+							expectedValue = uint64(bits.RotateLeft64(vs.x1, int(vs.x2)))
+						}
+					case *wazeroir.OperationRotr:
+						compileOperationFunc = func() {
+							err := compiler.compileRotr(o)
+							require.NoError(t, err)
+						}
+						is32Bit = o.Type == wazeroir.UnsignedInt32
+						if is32Bit {
+							expectedValue = uint64(bits.RotateLeft32(uint32(vs.x1), -int(vs.x2)))
+						} else {
+							expectedValue = uint64(bits.RotateLeft64(vs.x1, -int(vs.x2)))
+						}
+					}
+
+					// Setup the target values.
+					if is32Bit {
+						err := compiler.compileConstI32(&wazeroir.OperationConstI32{Value: uint32(vs.x1)})
+						require.NoError(t, err)
+						err = compiler.compileConstI32(&wazeroir.OperationConstI32{Value: uint32(vs.x2)})
+						require.NoError(t, err)
+					} else {
+						err := compiler.compileConstI64(&wazeroir.OperationConstI64{Value: vs.x1})
+						require.NoError(t, err)
+						err = compiler.compileConstI64(&wazeroir.OperationConstI64{Value: vs.x2})
+						require.NoError(t, err)
+					}
+
+					// Compile the operation.
+					compileOperationFunc()
+
+					// 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()
+					mem := newMemoryInst()
+					jitcall(
+						uintptr(unsafe.Pointer(&code[0])),
+						uintptr(unsafe.Pointer(eng)),
+						uintptr(unsafe.Pointer(&mem.Buffer[0])),
+					)
+
+					// Check the result.
+					require.Equal(t, uint64(1), eng.stackPointer)
+					if is32Bit {
+						require.Equal(t, uint32(expectedValue), uint32(eng.stack[eng.stackPointer-1]))
+					} else {
+						require.Equal(t, expectedValue, eng.stack[eng.stackPointer-1])
+					}
+				})
+			}
+		})
+	}
+}
+
 func TestAmd64Compiler_compileDiv(t *testing.T) {
 	t.Run("int32", func(t *testing.T) {
 		for _, signed := range []struct {