Skip to content

Commit

Permalink
jit(arm64): support for shl, shr, rotl, and rotr (tetratelabs#230)
Browse files Browse the repository at this point in the history 
Signed-off-by: r8d8 <ckryvomaz@gmail.com>
  • Loading branch information
@summerwind @r8d8
summerwind authored and r8d8 committed Feb 14, 2022
1 parent 28627d7 commit 5d205ea
Show file tree
Hide file tree
Showing 2 changed files with 216 additions and 5 deletions.
106 changes: 102 additions & 4 deletions wasm/jit/jit_arm64.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
115 changes: 114 additions & 1 deletion wasm/jit/jit_arm64_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import (
"context"
"fmt"
"math"
"math/bits"
"testing"
"unsafe"

Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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.stackTopAsUint32())
case wazeroir.UnsignedInt64:
require.Equal(t, x1<<(x2%64), env.stackTopAsUint64())
}
case wazeroir.OperationKindRotl:
switch unsignedInt {
case wazeroir.UnsignedInt32:
require.Equal(t, bits.RotateLeft32(uint32(x1), int(x2)), env.stackTopAsUint32())
case wazeroir.UnsignedInt64:
require.Equal(t, bits.RotateLeft64(x1, int(x2)), env.stackTopAsUint64())
}
case wazeroir.OperationKindRotr:
switch unsignedInt {
case wazeroir.UnsignedInt32:
require.Equal(t, bits.RotateLeft32(uint32(x1), -int(x2)), env.stackTopAsUint32())
case wazeroir.UnsignedInt64:
require.Equal(t, bits.RotateLeft64(x1, -int(x2)), env.stackTopAsUint64())
}
}
})
}
Expand All @@ -899,6 +930,88 @@ func TestArm64Compiler_compile_And_Or_Xor(t *testing.T) {
}
}

func TestArm64Compiler_compileShr(t *testing.T) {
kind := wazeroir.OperationKindShr
t.Run(kind.String(), func(t *testing.T) {
for _, signedInt := range []wazeroir.SignedInt{
wazeroir.SignedInt32,
wazeroir.SignedInt64,
wazeroir.SignedUint32,
wazeroir.SignedUint64,
} {
signedInt := signedInt
t.Run(signedInt.String(), func(t *testing.T) {
for _, values := range [][2]uint64{
{0, 0}, {0, 1}, {1, 0}, {1, 1},
{1 << 31, 1}, {1, 1 << 31}, {1 << 31, 1 << 31},
{1 << 63, 1}, {1, 1 << 63}, {1 << 63, 1 << 63},
} {
x1, x2 := values[0], values[1]
t.Run(fmt.Sprintf("x1=0x%x,x2=0x%x", x1, x2), func(t *testing.T) {
env := newJITEnvironment()
compiler := env.requireNewCompiler(t)
err := compiler.emitPreamble()
require.NoError(t, err)

// Emit consts operands.
for _, v := range []uint64{x1, x2} {
switch signedInt {
case wazeroir.SignedInt32:
err = compiler.compileConstI32(&wazeroir.OperationConstI32{Value: uint32(int32(v))})
case wazeroir.SignedInt64:
err = compiler.compileConstI64(&wazeroir.OperationConstI64{Value: v})
case wazeroir.SignedUint32:
err = compiler.compileConstI32(&wazeroir.OperationConstI32{Value: uint32(v)})
case wazeroir.SignedUint64:
err = compiler.compileConstI64(&wazeroir.OperationConstI64{Value: v})
}
require.NoError(t, err)
}

// At this point, two values exist.
require.Equal(t, uint64(2), compiler.locationStack.sp)

// Emit the operation.
err = compiler.compileShr(&wazeroir.OperationShr{Type: signedInt})
require.NoError(t, err)

// We consumed two values, but push the result back.
require.Equal(t, uint64(1), compiler.locationStack.sp)
resultLocation := compiler.locationStack.peek()
// Plus the result must be located on a register.
require.True(t, resultLocation.onRegister())
// Also, the result must have an appropriate register type.
require.Equal(t, generalPurposeRegisterTypeInt, resultLocation.regType)

// Release the value to the memory stack again to verify the operation.
compiler.releaseRegisterToStack(resultLocation)
compiler.returnFunction()

// Compile and execute the code under test.
code, _, _, err := compiler.compile()
require.NoError(t, err)
env.exec(code)

// Check the stack.
require.Equal(t, uint64(1), env.stackPointer())

switch signedInt {
case wazeroir.SignedInt32:
require.Equal(t, int32(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}
Expand Down

0 comments on commit 5d205ea

Please sign in to comment.