Repository: tetratelabs/wazero
Pull request #775 (2022-08-29 05:34:50)

Change OperationSwap to OperationSet
This commit changes `wazeroir.OperationSwap` to `OperationSet` which is the 
combination of OperationSwap and Drop in the previous implementation.

Previously, `OperationSwap` was always followed by `OperationDrop` on the swapped value on top.
Because of that, we had a redundant register allocation in Swap.

As a result, we use only one register in `OperationSet` which is a part of translations of
`local.tee` and `local.set`.

Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>


Patch:
diff --git a/internal/engine/compiler/compiler.go b/internal/engine/compiler/compiler.go
index 079c604e..e14113d9 100644
--- a/internal/engine/compiler/compiler.go
+++ b/internal/engine/compiler/compiler.go
@@ -25,8 +25,8 @@ type compiler interface {
 	compileLabel(o *wazeroir.OperationLabel) (skipThisLabel bool)
 	// compileUnreachable adds instruction to perform wazeroir.OperationUnreachable.
 	compileUnreachable() error
-	// compileSwap adds instruction to perform wazeroir.OperationSwap.
-	compileSwap(o *wazeroir.OperationSwap) error
+	// compileSet adds instruction to perform wazeroir.OperationSet.
+	compileSet(o *wazeroir.OperationSet) error
 	// compileGlobalGet adds instructions to perform wazeroir.OperationGlobalGet.
 	compileGlobalGet(o *wazeroir.OperationGlobalGet) error
 	// compileGlobalSet adds instructions to perform wazeroir.OperationGlobalSet.
diff --git a/internal/engine/compiler/engine.go b/internal/engine/compiler/engine.go
index 61cc4a77..44e83f3e 100644
--- a/internal/engine/compiler/engine.go
+++ b/internal/engine/compiler/engine.go
@@ -876,8 +876,8 @@ func compileWasmFunction(_ wasm.Features, ir *wazeroir.CompilationResult) (*code
 			err = compiler.compileSelect(o)
 		case *wazeroir.OperationPick:
 			err = compiler.compilePick(o)
-		case *wazeroir.OperationSwap:
-			err = compiler.compileSwap(o)
+		case *wazeroir.OperationSet:
+			err = compiler.compileSet(o)
 		case *wazeroir.OperationGlobalGet:
 			err = compiler.compileGlobalGet(o)
 		case *wazeroir.OperationGlobalSet:
diff --git a/internal/engine/compiler/impl_amd64.go b/internal/engine/compiler/impl_amd64.go
index ff53aa34..27db2efa 100644
--- a/internal/engine/compiler/impl_amd64.go
+++ b/internal/engine/compiler/impl_amd64.go
@@ -224,27 +224,28 @@ func (c *amd64Compiler) compileUnreachable() error {
 	return nil
 }
 
-// compileSwap implements compiler.compileSwap for the amd64 architecture.
-func (c *amd64Compiler) compileSwap(o *wazeroir.OperationSwap) error {
-	index := int(c.locationStack.sp) - 1 - o.Depth
-	var x1, x2 *runtimeValueLocation
+// compileSet implements compiler.compileSet for the amd64 architecture.
+func (c *amd64Compiler) compileSet(o *wazeroir.OperationSet) error {
+	setTargetIndex := int(c.locationStack.sp) - 1 - o.Depth
+
 	if o.IsTargetVector {
-		x1, x2 = c.locationStack.stack[c.locationStack.sp-2], c.locationStack.stack[index]
-	} else {
-		x1, x2 = c.locationStack.peek(), c.locationStack.stack[index]
+		_ = c.locationStack.pop() // ignore the higher 64-bits.
 	}
-
-	if err := c.compileEnsureOnRegister(x1); err != nil {
+	v := c.locationStack.pop()
+	if err := c.compileEnsureOnRegister(v); err != nil {
 		return err
 	}
-	if err := c.compileEnsureOnRegister(x2); err != nil {
-		return err
+
+	targetLocation := c.locationStack.stack[setTargetIndex]
+	if targetLocation.onRegister() {
+		// We no longer need the register previously used by the target location.
+		c.locationStack.markRegisterUnused(targetLocation.register)
 	}
 
-	x1.register, x2.register = x2.register, x1.register
+	reg := v.register
+	targetLocation.setRegister(reg)
 	if o.IsTargetVector {
-		x1, x2 = c.locationStack.peek(), c.locationStack.stack[index+1]
-		x1.register, x2.register = x2.register, x1.register
+		c.locationStack.stack[setTargetIndex+1].setRegister(reg)
 	}
 	return nil
 }
diff --git a/internal/engine/compiler/impl_arm64.go b/internal/engine/compiler/impl_arm64.go
index a346d8ba..60b76deb 100644
--- a/internal/engine/compiler/impl_arm64.go
+++ b/internal/engine/compiler/impl_arm64.go
@@ -453,28 +453,28 @@ func (c *arm64Compiler) compileUnreachable() error {
 	return nil
 }
 
-// compileSwap implements compiler.compileSwap for the arm64 architecture.
-func (c *arm64Compiler) compileSwap(o *wazeroir.OperationSwap) error {
-	yIndex := int(c.locationStack.sp) - 1 - o.Depth
-	var x, y *runtimeValueLocation
+// compileSet implements compiler.compileSet for the arm64 architecture.
+func (c *arm64Compiler) compileSet(o *wazeroir.OperationSet) error {
+	setTargetIndex := int(c.locationStack.sp) - 1 - o.Depth
+
 	if o.IsTargetVector {
-		x, y = c.locationStack.stack[c.locationStack.sp-2], c.locationStack.stack[yIndex]
-	} else {
-		x, y = c.locationStack.peek(), c.locationStack.stack[yIndex]
+		_ = c.locationStack.pop()
 	}
-
-	if err := c.compileEnsureOnRegister(x); err != nil {
+	v := c.locationStack.pop()
+	if err := c.compileEnsureOnRegister(v); err != nil {
 		return err
 	}
 
-	if err := c.compileEnsureOnRegister(y); err != nil {
-		return err
+	targetLocation := c.locationStack.stack[setTargetIndex]
+	if targetLocation.onRegister() {
+		// We no longer need the register previously used by the target location.
+		c.markRegisterUnused(targetLocation.register)
 	}
 
-	x.register, y.register = y.register, x.register
+	reg := v.register
+	targetLocation.setRegister(reg)
 	if o.IsTargetVector {
-		x, y = c.locationStack.peek(), c.locationStack.stack[yIndex+1]
-		x.register, y.register = y.register, x.register
+		c.locationStack.stack[setTargetIndex+1].setRegister(reg)
 	}
 	return nil
 }
diff --git a/internal/engine/interpreter/interpreter.go b/internal/engine/interpreter/interpreter.go
index bd2c7e74..2aefb068 100644
--- a/internal/engine/interpreter/interpreter.go
+++ b/internal/engine/interpreter/interpreter.go
@@ -390,7 +390,7 @@ func (e *engine) lowerIR(ir *wazeroir.CompilationResult) (*code, error) {
 			op.us = make([]uint64, 1)
 			op.us[0] = uint64(o.Depth)
 			op.b3 = o.IsTargetVector
-		case *wazeroir.OperationSwap:
+		case *wazeroir.OperationSet:
 			op.us = make([]uint64, 1)
 			op.us[0] = uint64(o.Depth)
 			op.b3 = o.IsTargetVector
@@ -934,14 +934,15 @@ func (ce *callEngine) callNativeFunc(ctx context.Context, callCtx *wasm.CallCont
 				ce.pushValue(ce.stack[index+1])
 			}
 			frame.pc++
-		case wazeroir.OperationKindSwap:
+		case wazeroir.OperationKindSet:
 			if op.b3 { // V128 value target.
 				lowIndex := len(ce.stack) - 1 - int(op.us[0])
-				ce.stack[len(ce.stack)-2], ce.stack[lowIndex] = ce.stack[lowIndex], ce.stack[len(ce.stack)-2]
-				ce.stack[len(ce.stack)-1], ce.stack[lowIndex+1] = ce.stack[lowIndex+1], ce.stack[len(ce.stack)-1]
+				highIndex := lowIndex + 1
+				hi, lo := ce.popValue(), ce.popValue()
+				ce.stack[lowIndex], ce.stack[highIndex] = lo, hi
 			} else {
 				index := len(ce.stack) - 1 - int(op.us[0])
-				ce.stack[len(ce.stack)-1], ce.stack[index] = ce.stack[index], ce.stack[len(ce.stack)-1]
+				ce.stack[index] = ce.popValue()
 			}
 			frame.pc++
 		case wazeroir.OperationKindGlobalGet:
diff --git a/internal/wazeroir/compiler.go b/internal/wazeroir/compiler.go
index f95de602..1737db12 100644
--- a/internal/wazeroir/compiler.go
+++ b/internal/wazeroir/compiler.go
@@ -826,15 +826,13 @@ operatorSwitch:
 			c.emit(
 				// +2 because we already popped the operands for this operation from the c.stack before
 				// called localDepth ^^,
-				&OperationSwap{Depth: depth + 2, IsTargetVector: isVector},
-				&OperationDrop{Depth: &InclusiveRange{Start: 0, End: 1}},
+				&OperationSet{Depth: depth + 2, IsTargetVector: isVector},
 			)
 		} else {
 			c.emit(
 				// +1 because we already popped the operands for this operation from the c.stack before
 				// called localDepth ^^,
-				&OperationSwap{Depth: depth + 1, IsTargetVector: isVector},
-				&OperationDrop{Depth: &InclusiveRange{Start: 0, End: 0}},
+				&OperationSet{Depth: depth + 1, IsTargetVector: isVector},
 			)
 		}
 	case wasm.OpcodeLocalTee:
@@ -847,14 +845,12 @@ operatorSwitch:
 		if isVector {
 			c.emit(
 				&OperationPick{Depth: 1, IsTargetVector: isVector},
-				&OperationSwap{Depth: depth + 2, IsTargetVector: isVector},
-				&OperationDrop{Depth: &InclusiveRange{Start: 0, End: 1}},
+				&OperationSet{Depth: depth + 2, IsTargetVector: isVector},
 			)
 		} else {
 			c.emit(
 				&OperationPick{Depth: 0, IsTargetVector: isVector},
-				&OperationSwap{Depth: depth + 1, IsTargetVector: isVector},
-				&OperationDrop{Depth: &InclusiveRange{Start: 0, End: 0}},
+				&OperationSet{Depth: depth + 1, IsTargetVector: isVector},
 			)
 		}
 	case wasm.OpcodeGlobalGet:
diff --git a/internal/wazeroir/format.go b/internal/wazeroir/format.go
index ae8223b1..17fb5d79 100644
--- a/internal/wazeroir/format.go
+++ b/internal/wazeroir/format.go
@@ -49,7 +49,7 @@ func formatOperation(w io.StringWriter, b Operation) {
 		str = "select"
 	case *OperationPick:
 		str = fmt.Sprintf("pick %d (is_vector=%v)", o.Depth, o.IsTargetVector)
-	case *OperationSwap:
+	case *OperationSet:
 		str = fmt.Sprintf("swap %d (is_vector=%v)", o.Depth, o.IsTargetVector)
 	case *OperationGlobalGet:
 		str = fmt.Sprintf("global.get %d", o.Index)
diff --git a/internal/wazeroir/operations.go b/internal/wazeroir/operations.go
index 22503a9b..abc9bfc9 100644
--- a/internal/wazeroir/operations.go
+++ b/internal/wazeroir/operations.go
@@ -160,7 +160,7 @@ func (o OperationKind) String() (ret string) {
 		ret = "Select"
 	case OperationKindPick:
 		ret = "Pick"
-	case OperationKindSwap:
+	case OperationKindSet:
 		ret = "Swap"
 	case OperationKindGlobalGet:
 		ret = "GlobalGet"
@@ -443,8 +443,8 @@ const (
 	OperationKindSelect
 	// OperationKindPick is the kind for OperationPick.
 	OperationKindPick
-	// OperationKindSwap is the kind for OperationSwap.
-	OperationKindSwap
+	// OperationKindSet is the kind for OperationSet.
+	OperationKindSet
 	// OperationKindGlobalGet is the kind for OperationGlobalGet.
 	OperationKindGlobalGet
 	// OperationKindGlobalSet is the kind for OperationGlobalSet.
@@ -957,20 +957,20 @@ func (*OperationPick) Kind() OperationKind {
 	return OperationKindPick
 }
 
-// OperationSwap implements Operation.
+// OperationSet implements Operation.
 //
-// The engines are expected to swap the top value of the stack and the one specified by
-// OperationSwap.Depth.
-type OperationSwap struct {
-	// Depth is the location of the pick target in the uint64 value stack at runtime.
+// The engines are expected to set the top value of the stack to the location specified by
+// OperationSet.Depth.
+type OperationSet struct {
+	// Depth is the location of the set target in the uint64 value stack at runtime.
 	// If IsTargetVector=true, this points the location of the lower 64-bits of the vector.
 	Depth          int
 	IsTargetVector bool
 }
 
 // Kind implements Operation.Kind
-func (*OperationSwap) Kind() OperationKind {
-	return OperationKindSwap
+func (*OperationSet) Kind() OperationKind {
+	return OperationKindSet
 }
 
 // OperationGlobalGet implements Operation.


Test patch:
diff --git a/internal/engine/compiler/compiler_stack_test.go b/internal/engine/compiler/compiler_stack_test.go
index be93bc75..6f9a856e 100644
--- a/internal/engine/compiler/compiler_stack_test.go
+++ b/internal/engine/compiler/compiler_stack_test.go
@@ -648,7 +648,7 @@ func TestCompiler_compileSwap_v128(t *testing.T) {
 			}
 
 			// Swap x1 and x2.
-			err = compiler.compileSwap(&wazeroir.OperationSwap{Depth: 4, IsTargetVector: true})
+			err = compiler.compileSet(&wazeroir.OperationSet{Depth: 4, IsTargetVector: true})
 			require.NoError(t, err)
 
 			require.NoError(t, compiler.compileReturnFunction())
@@ -661,18 +661,16 @@ func TestCompiler_compileSwap_v128(t *testing.T) {
 			env.exec(code)
 
 			require.Equal(t, nativeCallStatusCodeReturned, env.compilerStatus())
-			require.Equal(t, uint64(5), env.stackPointer())
+			require.Equal(t, uint64(3), env.stackPointer())
 
 			st := env.stack()
 			require.Equal(t, x2Lo, st[0])
 			require.Equal(t, x2Hi, st[1])
-			require.Equal(t, x1Lo, st[3])
-			require.Equal(t, x1Hi, st[4])
 		})
 	}
 }
 
-func TestCompiler_compileSwap(t *testing.T) {
+func TestCompiler_compileSet(t *testing.T) {
 	var x1Value, x2Value int64 = 100, 200
 	tests := []struct {
 		x1OnConditionalRegister, x1OnRegister, x2OnRegister bool
@@ -721,7 +719,7 @@ func TestCompiler_compileSwap(t *testing.T) {
 			}
 
 			// Swap x1 and x2.
-			err = compiler.compileSwap(&wazeroir.OperationSwap{Depth: 2})
+			err = compiler.compileSet(&wazeroir.OperationSet{Depth: 2})
 			require.NoError(t, err)
 
 			require.NoError(t, compiler.compileReturnFunction())
@@ -733,10 +731,9 @@ func TestCompiler_compileSwap(t *testing.T) {
 			// Run code.
 			env.exec(code)
 
-			require.Equal(t, uint64(3), env.stackPointer())
-			// Check values are swapped.
+			require.Equal(t, uint64(2), env.stackPointer())
+			// Check the value was set.
 			require.Equal(t, uint64(x1Value), env.stack()[0])
-			require.Equal(t, uint64(x2Value), env.stack()[2])
 		})
 	}
 }
diff --git a/internal/wazeroir/compiler_test.go b/internal/wazeroir/compiler_test.go
index bb2d2781..f0315219 100644
--- a/internal/wazeroir/compiler_test.go
+++ b/internal/wazeroir/compiler_test.go
@@ -1078,10 +1078,8 @@ func TestCompile_Locals(t *testing.T) {
 			expected: []Operation{
 				// [p[0].lo, p[1].hi] -> [p[0].lo, p[1].hi, 0x01, 0x02]
 				&OperationV128Const{Lo: 0x01, Hi: 0x02},
-				// [p[0].lo, p[1].hi, 0x01, 0x02] -> [0x01, 0x02, p[0].lo, p[1].hi]
-				&OperationSwap{Depth: 3, IsTargetVector: true},
-				// [0x01, 0x02, p[0].lo, p[1].hi] -> [0x02, 0x01]
-				&OperationDrop{Depth: &InclusiveRange{Start: 0, End: 1}},
+				// [p[0].lo, p[1].hi, 0x01, 0x02] -> [0x01, 0x02]
+				&OperationSet{Depth: 3, IsTargetVector: true},
 				&OperationDrop{Depth: &InclusiveRange{Start: 0, End: 1}},
 				&OperationBr{Target: &BranchTarget{}}, // return!
 			},
@@ -1099,8 +1097,7 @@ func TestCompile_Locals(t *testing.T) {
 			},
 			expected: []Operation{
 				&OperationConstI32{Value: 0x1},
-				&OperationSwap{Depth: 1, IsTargetVector: false},
-				&OperationDrop{Depth: &InclusiveRange{Start: 0, End: 0}},
+				&OperationSet{Depth: 1, IsTargetVector: false},
 				&OperationDrop{Depth: &InclusiveRange{Start: 0, End: 0}},
 				&OperationBr{Target: &BranchTarget{}}, // return!
 			},
@@ -1125,10 +1122,8 @@ func TestCompile_Locals(t *testing.T) {
 				&OperationV128Const{Lo: 0, Hi: 0},
 				// [p[0].lo, p[1].hi] -> [p[0].lo, p[1].hi, 0x01, 0x02]
 				&OperationV128Const{Lo: 0x01, Hi: 0x02},
-				// [p[0].lo, p[1].hi, 0x01, 0x02] -> [0x01, 0x02, p[0].lo, p[1].hi]
-				&OperationSwap{Depth: 3, IsTargetVector: true},
-				// [p[0].lo, 0x02, 0x01, p[1].hi] -> [0x02, 0x01]
-				&OperationDrop{Depth: &InclusiveRange{Start: 0, End: 1}},
+				// [p[0].lo, p[1].hi, 0x01, 0x02] -> [0x01, 0x02]
+				&OperationSet{Depth: 3, IsTargetVector: true},
 				&OperationDrop{Depth: &InclusiveRange{Start: 0, End: 1}},
 				&OperationBr{Target: &BranchTarget{}}, // return!
 			},
@@ -1151,10 +1146,8 @@ func TestCompile_Locals(t *testing.T) {
 				&OperationV128Const{Lo: 0x01, Hi: 0x02},
 				// [p[0].lo, p[1].hi, 0x01, 0x02] -> [p[0].lo, p[1].hi, 0x01, 0x02, 0x01, 0x02]
 				&OperationPick{Depth: 1, IsTargetVector: true},
-				// [p[0].lo, p[1].hi, 0x01, 0x02, 0x01, 0x02] -> [0x01, 0x02, 0x01, 0x02, p[0].lo, p[1].hi]
-				&OperationSwap{Depth: 5, IsTargetVector: true},
-				// [0x01, 0x02, 0x01, 0x02, p[0].lo, p[1].hi] ->  [0x01, 0x02, 0x01, 0x02]
-				&OperationDrop{Depth: &InclusiveRange{Start: 0, End: 1}},
+				// [p[0].lo, p[1].hi, 0x01, 0x02, 0x01, 0x02] -> [0x01, 0x02, 0x01, 0x02]
+				&OperationSet{Depth: 5, IsTargetVector: true},
 				&OperationDrop{Depth: &InclusiveRange{Start: 0, End: 3}},
 				&OperationBr{Target: &BranchTarget{}}, // return!
 			},
@@ -1173,8 +1166,7 @@ func TestCompile_Locals(t *testing.T) {
 			expected: []Operation{
 				&OperationConstF32{math.Float32frombits(1)},
 				&OperationPick{Depth: 0, IsTargetVector: false},
-				&OperationSwap{Depth: 2, IsTargetVector: false},
-				&OperationDrop{Depth: &InclusiveRange{Start: 0, End: 0}},
+				&OperationSet{Depth: 2, IsTargetVector: false},
 				&OperationDrop{Depth: &InclusiveRange{Start: 0, End: 1}},
 				&OperationBr{Target: &BranchTarget{}}, // return!
 			},
@@ -1201,10 +1193,8 @@ func TestCompile_Locals(t *testing.T) {
 				&OperationV128Const{Lo: 0x01, Hi: 0x02},
 				// [p[0].lo, p[1].hi, 0x01, 0x02] -> [p[0].lo, p[1].hi, 0x01, 0x02, 0x01, 0x02]
 				&OperationPick{Depth: 1, IsTargetVector: true},
-				// [p[0].lo, p[1].hi, 0x01, 0x02, 0x01, 0x2] -> [0x01, 0x02, 0x01, 0x02, p[0].lo, p[1].hi]
-				&OperationSwap{Depth: 5, IsTargetVector: true},
-				// [0x01, 0x02, 0x01, 0x02, p[0].lo, p[1].hi] ->  [0x01, 0x02, 0x01, 0x02]
-				&OperationDrop{Depth: &InclusiveRange{Start: 0, End: 1}},
+				// [p[0].lo, p[1].hi, 0x01, 0x02, 0x01, 0x2] -> [0x01, 0x02, 0x01, 0x02]
+				&OperationSet{Depth: 5, IsTargetVector: true},
 				&OperationDrop{Depth: &InclusiveRange{Start: 0, End: 3}},
 				&OperationBr{Target: &BranchTarget{}}, // return!
 			},