Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

jit(arm64): support for shl, shr, rotl, and rotr #230

Merged
merged 2 commits into from
Feb 13, 2022
Merged

jit(arm64): support for shl, shr, rotl, and rotr #230

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c1156b7
jit(arm64): support for shl, shr, rotl, and rotr
summerwind Feb 13, 2022
feb5ed2
Update wasm/jit/jit_arm64.go
mathetake Feb 13, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
106 changes: 102 additions & 4 deletions wasm/jit/jit_arm64.go
View file
Open in desktop
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 {
codefromthecrypt marked this conversation as resolved.
Show resolved Hide resolved
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.
mathetake marked this conversation as resolved.
Show resolved Hide resolved
c.applyRegisterToRegisterInstruction(neginst, x2.register, x2.register)
Copy link
Contributor Author

@summerwind summerwind Feb 13, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I adopted this way because I checked Go Assembly output of the bits.RotateLeft() sample code described in the official documentation, and it looked like this:

...
0x0258 00600 (main.go:11)       MOVD    $15, R0
0x025c 00604 (main.go:11)       MOVD    $2, R1
0x0260 00608 (main.go:11)       RORW    R1, R0, R0
...
0x0130 00304 (main.go:10)       MOVD    $15, R0
0x0134 00308 (main.go:10)       MOVD    $-2, R1
0x0138 00312 (main.go:10)       RORW    R1, R0, R1
...

I think this is valid, but I'm not sure this is the best way.
I'd love to hear your comments!

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ok confirmed that NEG is actually used in register allocated case

func fooooo(x uint32, k int) uint32 {
	return bits.RotateLeft32(x, k)
}

is compiled to

"".fooooo STEXT size=32 args=0x18 locals=0x0 funcid=0x0 leaf
	0x0000 00000 (main.go:12)	TEXT	"".fooooo(SB), LEAF|NOFRAME|ABIInternal, $0-24
	0x0000 00000 (main.go:12)	FUNCDATA	ZR, gclocals·33cdeccccebe80329f1fdbee7f5874cb(SB)
	0x0000 00000 (main.go:12)	FUNCDATA	$1, gclocals·33cdeccccebe80329f1fdbee7f5874cb(SB)
	0x0000 00000 (main.go:12)	FUNCDATA	$5, "".fooooo.arginfo1(SB)
	0x0000 00000 (main.go:13)	MOVD	"".k+8(FP), R0
	0x0004 00004 (main.go:13)	NEG	R0, R0
	0x0008 00008 (main.go:13)	MOVWU	"".x(FP), R1
	0x000c 00012 (main.go:13)	RORW	R0, R1, R0
	0x0010 00016 (main.go:13)	MOVW	R0, "".~r2+16(FP)
	0x0014 00020 (main.go:13)	RET	(R30)
	0x0000 e0 0b 40 f9 e0 03 00 cb e1 0b 40 b9 20 2c c0 1a  ..@.......@. ,..
	0x0010 e0 1b 00 b9 c0 03 5f d6 00 00 00 00 00 00 00 00  ......_.........

Thank you for your efforts!


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
Open in desktop
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