diff --git a/wasm/jit/jit_arm64.go b/wasm/jit/jit_arm64.go index 6ca14dcbff..0c1cd0fca0 100644 --- a/wasm/jit/jit_arm64.go +++ b/wasm/jit/jit_arm64.go @@ -861,20 +861,118 @@ func (c *arm64Compiler) compileXor(o *wazeroir.OperationXor) error { return nil } +// compileShl implements compiler.compileShl for the arm64 architecture. func (c *arm64Compiler) compileShl(o *wazeroir.OperationShl) error { - return fmt.Errorf("TODO: unsupported on arm64") + x1, x2, err := c.popTwoValuesOnRegisters() + if err != nil { + return err + } + + if isZeroRegister(x1.register) || isZeroRegister(x2.register) { + c.locationStack.pushValueLocationOnRegister(x1.register) + return nil + } + + var inst obj.As + switch o.Type { + case wazeroir.UnsignedInt32: + inst = arm64.ALSLW + case wazeroir.UnsignedInt64: + inst = arm64.ALSL + } + + c.applyTwoRegistersToRegisterInstruction(inst, x2.register, x1.register, x1.register) + c.locationStack.pushValueLocationOnRegister(x1.register) + return nil } +// compileShr implements compiler.compileShr for the arm64 architecture. func (c *arm64Compiler) compileShr(o *wazeroir.OperationShr) error { - return fmt.Errorf("TODO: unsupported on arm64") + x1, x2, err := c.popTwoValuesOnRegisters() + if err != nil { + return err + } + + if isZeroRegister(x1.register) || isZeroRegister(x2.register) { + c.locationStack.pushValueLocationOnRegister(x1.register) + return nil + } + + var inst obj.As + switch o.Type { + case wazeroir.SignedInt32: + inst = arm64.AASRW + case wazeroir.SignedInt64: + inst = arm64.AASR + case wazeroir.SignedUint32: + inst = arm64.ALSRW + case wazeroir.SignedUint64: + inst = arm64.ALSR + } + + c.applyTwoRegistersToRegisterInstruction(inst, x2.register, x1.register, x1.register) + c.locationStack.pushValueLocationOnRegister(x1.register) + return nil } +// compileRotl implements compiler.compileRotl for the arm64 architecture. func (c *arm64Compiler) compileRotl(o *wazeroir.OperationRotl) error { - return fmt.Errorf("TODO: unsupported on arm64") + x1, x2, err := c.popTwoValuesOnRegisters() + if err != nil { + return err + } + + if isZeroRegister(x1.register) || isZeroRegister(x2.register) { + c.locationStack.pushValueLocationOnRegister(x1.register) + return nil + } + + var ( + inst obj.As + neginst obj.As + ) + + switch o.Type { + case wazeroir.UnsignedInt32: + inst = arm64.ARORW + neginst = arm64.ANEGW + case wazeroir.UnsignedInt64: + inst = arm64.AROR + neginst = arm64.ANEG + } + + // Arm64 doesn't have rotate left instruction. + // The shift amount needs to be converted to a negative number, similar to assembly output of bits.RotateLeft. + c.applyRegisterToRegisterInstruction(neginst, x2.register, x2.register) + + c.applyTwoRegistersToRegisterInstruction(inst, x2.register, x1.register, x1.register) + c.locationStack.pushValueLocationOnRegister(x1.register) + return nil } +// compileRotr implements compiler.compileRotr for the arm64 architecture. func (c *arm64Compiler) compileRotr(o *wazeroir.OperationRotr) error { - return fmt.Errorf("TODO: unsupported on arm64") + x1, x2, err := c.popTwoValuesOnRegisters() + if err != nil { + return err + } + + if isZeroRegister(x1.register) || isZeroRegister(x2.register) { + c.locationStack.pushValueLocationOnRegister(x1.register) + return nil + } + + var inst obj.As + switch o.Type { + case wazeroir.UnsignedInt32: + inst = arm64.ARORW + case wazeroir.UnsignedInt64: + inst = arm64.AROR + } + + c.applyTwoRegistersToRegisterInstruction(inst, x2.register, x1.register, x1.register) + c.locationStack.pushValueLocationOnRegister(x1.register) + return nil } func (c *arm64Compiler) compileAbs(o *wazeroir.OperationAbs) error { diff --git a/wasm/jit/jit_arm64_test.go b/wasm/jit/jit_arm64_test.go index 5e7f7465a4..fa291f3e4c 100644 --- a/wasm/jit/jit_arm64_test.go +++ b/wasm/jit/jit_arm64_test.go @@ -7,6 +7,7 @@ import ( "context" "fmt" "math" + "math/bits" "testing" "unsafe" @@ -797,11 +798,14 @@ func TestArm64Compiler_compile_Add_Sub_Mul(t *testing.T) { } } -func TestArm64Compiler_compile_And_Or_Xor(t *testing.T) { +func TestArm64Compiler_compile_And_Or_Xor_Shl_Rotr(t *testing.T) { for _, kind := range []wazeroir.OperationKind{ wazeroir.OperationKindAnd, wazeroir.OperationKindOr, wazeroir.OperationKindXor, + wazeroir.OperationKindShl, + wazeroir.OperationKindRotl, + wazeroir.OperationKindRotr, } { kind := kind t.Run(kind.String(), func(t *testing.T) { @@ -845,6 +849,12 @@ func TestArm64Compiler_compile_And_Or_Xor(t *testing.T) { err = compiler.compileOr(&wazeroir.OperationOr{Type: unsignedInt}) case wazeroir.OperationKindXor: err = compiler.compileXor(&wazeroir.OperationXor{Type: unsignedInt}) + case wazeroir.OperationKindShl: + err = compiler.compileShl(&wazeroir.OperationShl{Type: unsignedInt}) + case wazeroir.OperationKindRotl: + err = compiler.compileRotl(&wazeroir.OperationRotl{Type: unsignedInt}) + case wazeroir.OperationKindRotr: + err = compiler.compileRotr(&wazeroir.OperationRotr{Type: unsignedInt}) } require.NoError(t, err) @@ -890,6 +900,27 @@ func TestArm64Compiler_compile_And_Or_Xor(t *testing.T) { case wazeroir.UnsignedInt64: require.Equal(t, x1^x2, env.stackTopAsUint64()) } + case wazeroir.OperationKindShl: + switch unsignedInt { + case wazeroir.UnsignedInt32: + require.Equal(t, uint32(x1)<>(uint32(x2)%32), env.stackTopAsInt32()) + case wazeroir.SignedInt64: + require.Equal(t, int64(x1)>>(x2%64), env.stackTopAsInt64()) + case wazeroir.SignedUint32: + require.Equal(t, uint32(x1)>>(uint32(x2)%32), env.stackTopAsUint32()) + case wazeroir.SignedUint64: + require.Equal(t, x1>>(x2%64), env.stackTopAsUint64()) + } + }) + } + }) + } + }) +} + func TestArm64Compiler_compielePick(t *testing.T) { const pickTargetValue uint64 = 12345 op := &wazeroir.OperationPick{Depth: 1}