Repository: tetratelabs/wazero
Pull request #73 (2021-12-12 09:19:41)

jit: support Swap instruction
part of #65 

Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>


Patch:
diff --git a/wasm/jit/jit_amd64.go b/wasm/jit/jit_amd64.go
index 287665bb..bd9ff4dc 100644
--- a/wasm/jit/jit_amd64.go
+++ b/wasm/jit/jit_amd64.go
@@ -92,7 +92,9 @@ func (e *engine) compileWasmFunction(f *wasm.FunctionInstance) (*compiledWasmFun
 				return nil, fmt.Errorf("error handling pick operation: %w", err)
 			}
 		case *wazeroir.OperationSwap:
-			return nil, fmt.Errorf("unsupported operation in JIT compiler: %v", o)
+			if err := builder.handleSwap(o); err != nil {
+				return nil, fmt.Errorf("error handling swap operation: %w", err)
+			}
 		case *wazeroir.OperationGlobalGet:
 			return nil, fmt.Errorf("unsupported operation in JIT compiler: %v", o)
 		case *wazeroir.OperationGlobalSet:
@@ -308,6 +310,91 @@ func (b *amd64Builder) handleUnreachable() {
 	b.returnFunction()
 }
 
+func (b *amd64Builder) handleSwap(o *wazeroir.OperationSwap) error {
+	index := len(b.locationStack.stack) - 1 - o.Depth
+	// Note that, in theory, the register types and value types
+	// are the same between these swap targets as swap operations
+	// are generated from local.set,tee instructions in Wasm.
+	x1 := b.locationStack.stack[len(b.locationStack.stack)-1]
+	x2 := b.locationStack.stack[index]
+
+	// If x1 is on the conditional register, we must move it to a gp
+	// register before swap.
+	if x1.onConditionalRegister() {
+		if err := b.moveConditionalToGPRegister(x1); err != nil {
+			return err
+		}
+	}
+
+	if x1.onRegister() && x2.onRegister() {
+		x1.register, x2.register = x2.register, x1.register
+	} else if x1.onRegister() && x2.onStack() {
+		reg := x1.register
+		// Save x1's value to the temporary top of the stack.
+		tmpStackLocation := b.locationStack.pushValueOnRegister(reg)
+		b.releaseRegisterToStack(tmpStackLocation)
+		// Then move the x2's value to the x1's register location.
+		x2.register = reg
+		b.moveStackToRegister(x2)
+		// Now move the x1's value to the x1's stack location.
+		b.releaseRegisterToStack(x1)
+		// Next we move the saved x1's value to the register.
+		tmpStackLocation.setRegister(reg)
+		b.moveStackToRegister(tmpStackLocation)
+		// Finally move the x1's value in the register to the x2's stack location.
+		b.locationStack.releaseRegister(x1)
+		b.locationStack.releaseRegister(tmpStackLocation)
+		x2.setRegister(reg)
+		b.locationStack.markRegisterUsed(reg)
+		_ = b.locationStack.pop() // Delete tmpStackLocation.
+	} else if x1.onStack() && x2.onRegister() {
+		reg := x2.register
+		// Save x2's value to the temporary top of the stack.
+		tmpStackLocation := b.locationStack.pushValueOnRegister(reg)
+		b.releaseRegisterToStack(tmpStackLocation)
+		// Then move the x1's value to the x2's register location.
+		x1.register = reg
+		b.moveStackToRegister(x1)
+		// Now move the x1's value to the x2's stack location.
+		b.releaseRegisterToStack(x2)
+		// Next we move the saved x2's value to the register.
+		tmpStackLocation.setRegister(reg)
+		b.moveStackToRegister(tmpStackLocation)
+		// Finally move the x2's value in the register to the x2's stack location.
+		b.locationStack.releaseRegister(x2)
+		b.locationStack.releaseRegister(tmpStackLocation)
+		x1.setRegister(reg)
+		b.locationStack.markRegisterUsed(reg)
+		_ = b.locationStack.pop() // Delete tmpStackLocation.
+	} else if x1.onStack() && x2.onStack() {
+		reg, err := b.allocateRegister(x1.registerType())
+		if err != nil {
+			return err
+		}
+		// First we move the x2's value to the temp register.
+		x2.setRegister(reg)
+		b.moveStackToRegister(x2)
+		// Save x2's value to the temporary top of the stack.
+		tmpStackLocation := b.locationStack.pushValueOnRegister(reg)
+		b.releaseRegisterToStack(tmpStackLocation)
+		// Then move the x1's value to the x2's register location.
+		x1.register = reg
+		b.moveStackToRegister(x1)
+		// Now move the x1's value to the x2's stack location.
+		b.releaseRegisterToStack(x2)
+		// Next we move the saved x2's value to the register.
+		tmpStackLocation.setRegister(reg)
+		b.moveStackToRegister(tmpStackLocation)
+		// Finally move the x2's value in the register to the x2's stack location.
+		b.locationStack.releaseRegister(x2)
+		b.locationStack.releaseRegister(tmpStackLocation)
+		x1.setRegister(reg)
+		b.locationStack.markRegisterUsed(reg)
+		_ = b.locationStack.pop() // Delete tmpStackLocation.
+	}
+	return nil
+}
+
 func (b *amd64Builder) handleBr(o *wazeroir.OperationBr) error {
 	if o.Target.IsReturnTarget() {
 		// Release all the registers as our calling convention requires the callee-save.
diff --git a/wasm/jit/jit_value_location_amd64.go b/wasm/jit/jit_value_location_amd64.go
index aeef7b24..711b5fb9 100644
--- a/wasm/jit/jit_value_location_amd64.go
+++ b/wasm/jit/jit_value_location_amd64.go
@@ -194,6 +194,7 @@ func (s *valueLocationStack) peek() (loc *valueLocation) {
 func (s *valueLocationStack) releaseRegister(loc *valueLocation) {
 	delete(s.usedRegisters, loc.register)
 	loc.register = -1
+	loc.conditionalRegister = conditionalRegisterStateUnset
 }
 
 func (s *valueLocationStack) markRegisterUsed(reg int16) {


Test patch:
diff --git a/wasm/jit/jit_amd64_test.go b/wasm/jit/jit_amd64_test.go
index 9b833e1e..c32a8e14 100644
--- a/wasm/jit/jit_amd64_test.go
+++ b/wasm/jit/jit_amd64_test.go
@@ -1612,3 +1612,73 @@ func TestAmd64Builder_handleSelect(t *testing.T) {
 		})
 	}
 }
+
+func TestAmd64Builder_handleSwap(t *testing.T) {
+	var x1Value, x2Value int64 = 100, 200
+	for i, tc := range []struct {
+		x1OnConditionalRegister, x1OnRegister, x2OnRegister bool
+	}{
+		{x1OnRegister: true, x2OnRegister: true},
+		{x1OnRegister: true, x2OnRegister: false},
+		{x1OnRegister: false, x2OnRegister: true},
+		{x1OnRegister: false, x2OnRegister: false},
+		// x1 on conditional register
+		{x1OnConditionalRegister: true, x2OnRegister: false},
+		{x1OnConditionalRegister: true, x2OnRegister: true},
+	} {
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+			eng := newEngine()
+			builder := requireNewBuilder(t)
+			builder.initializeReservedRegisters()
+			if tc.x2OnRegister {
+				x2 := builder.locationStack.pushValueOnRegister(x86.REG_R10)
+				builder.movConstToRegister(x2Value, x2.register)
+			} else {
+				x2 := builder.locationStack.pushValueOnStack()
+				eng.stack[x2.stackPointer] = uint64(x2Value)
+			}
+			_ = builder.locationStack.pushValueOnStack() // Dummy value!
+			if tc.x1OnRegister && !tc.x1OnConditionalRegister {
+				x1 := builder.locationStack.pushValueOnRegister(x86.REG_AX)
+				builder.movConstToRegister(x1Value, x1.register)
+			} else if !tc.x1OnConditionalRegister {
+				x1 := builder.locationStack.pushValueOnStack()
+				eng.stack[x1.stackPointer] = uint64(x1Value)
+			} else {
+				builder.movConstToRegister(0, x86.REG_AX)
+				cmp := builder.newProg()
+				cmp.As = x86.ACMPQ
+				cmp.From.Type = obj.TYPE_REG
+				cmp.From.Reg = x86.REG_AX
+				cmp.To.Type = obj.TYPE_CONST
+				cmp.To.Offset = 0
+				cmp.To.Offset = 0
+				builder.addInstruction(cmp)
+				builder.locationStack.pushValueOnConditionalRegister(conditionalRegisterStateE)
+				x1Value = 1
+			}
+
+			// Swap x1 and x2.
+			err := builder.handleSwap(&wazeroir.OperationSwap{Depth: 2})
+			require.NoError(t, err)
+			// To verify the behavior, we release all the registers to stack locations.
+			builder.releaseAllRegistersToStack()
+			builder.returnFunction()
+
+			// Assemble.
+			code, err := builder.assemble()
+			require.NoError(t, err)
+			// Run code.
+			mem := newMemoryInst()
+			jitcall(
+				uintptr(unsafe.Pointer(&code[0])),
+				uintptr(unsafe.Pointer(eng)),
+				uintptr(unsafe.Pointer(&mem.Buffer[0])),
+			)
+			require.Equal(t, uint64(3), eng.currentStackPointer)
+			// Check values are swapped.
+			require.Equal(t, uint64(x1Value), eng.stack[0])
+			require.Equal(t, uint64(x2Value), eng.stack[2])
+		})
+	}
+}