cmd/compile: add wasm stack optimization

Go's SSA instructions only operate on registers. For example, an add
instruction would read two registers, do the addition and then write
to a register. WebAssembly's instructions, on the other hand, operate
on the stack. The add instruction first pops two values from the stack,
does the addition, then pushes the result to the stack. To fulfill
Go's semantics, one needs to map Go's single add instruction to
4 WebAssembly instructions:
- Push the value of local variable A to the stack
- Push the value of local variable B to the stack
- Do addition
- Write value from stack to local variable C

Now consider that B was set to the constant 42 before the addition:
- Push constant 42 to the stack
- Write value from stack to local variable B

This works, but is inefficient. Instead, the stack is used directly
by inlining instructions if possible. With inlining it becomes:
- Push the value of local variable A to the stack (add)
- Push constant 42 to the stack (constant)
- Do addition (add)
- Write value from stack to local variable C (add)

Note that the two SSA instructions can not be generated sequentially
anymore, because their WebAssembly instructions are interleaved.

Design doc: https://docs.google.com/document/d/131vjr4DH6JFnb-blm_uRdaC0_Nv3OUwjEY5qVCxCup4

Updates #18892

Change-Id: Ie35e1c0bebf4985fddda0d6330eb2066f9ad6dec
Reviewed-on: https://go-review.googlesource.com/103535
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Cherry Zhang <cherryyz@google.com>

1 files changed, 58 insertions, 32 deletions

diff --git a/src/cmd/compile/internal/wasm/ssa.go b/src/cmd/compile/internal/wasm/ssa.go
index d4bd7fbdec..8daf749a12 100644
--- a/src/cmd/compile/internal/wasm/ssa.go
+++ b/src/cmd/compile/internal/wasm/ssa.go
@@ -78,7 +78,7 @@ func ssaGenBlock(s *gc.SSAGenState, b, next *ssa.Block) {
 		goToBlock(b.Succs[0].Block(), true)
 
 	case ssa.BlockIf:
-		getReg32(s, b.Control)
+		getValue32(s, b.Control)
 		s.Prog(wasm.AI32Eqz)
 		s.Prog(wasm.AIf)
 		goToBlock(b.Succs[1].Block(), false)
@@ -113,6 +113,10 @@ func ssaGenBlock(s *gc.SSAGenState, b, next *ssa.Block) {
 
 	// Entry point for the next block. Used by the JMP in goToBlock.
 	s.Prog(wasm.ARESUMEPOINT)
+
+	if s.OnWasmStackSkipped != 0 {
+		panic("wasm: bad stack")
+	}
 }
 
 func ssaGenValue(s *gc.SSAGenState, v *ssa.Value) {
@@ -124,33 +128,33 @@ func ssaGenValue(s *gc.SSAGenState, v *ssa.Value) {
 			s.Prog(wasm.ARESUMEPOINT)
 		}
 		if v.Op == ssa.OpWasmLoweredClosureCall {
-			getReg64(s, v.Args[1])
+			getValue64(s, v.Args[1])
 			setReg(s, wasm.REG_CTXT)
 		}
 		if sym, ok := v.Aux.(*obj.LSym); ok {
 			p := s.Prog(obj.ACALL)
 			p.To = obj.Addr{Type: obj.TYPE_MEM, Name: obj.NAME_EXTERN, Sym: sym}
 		} else {
-			getReg64(s, v.Args[0])
+			getValue64(s, v.Args[0])
 			p := s.Prog(obj.ACALL)
 			p.To = obj.Addr{Type: obj.TYPE_NONE}
 		}
 
 	case ssa.OpWasmLoweredMove:
-		getReg32(s, v.Args[0])
-		getReg32(s, v.Args[1])
+		getValue32(s, v.Args[0])
+		getValue32(s, v.Args[1])
 		i32Const(s, int32(v.AuxInt))
 		p := s.Prog(wasm.ACall)
 		p.To = obj.Addr{Type: obj.TYPE_MEM, Name: obj.NAME_EXTERN, Sym: gc.WasmMove}
 
 	case ssa.OpWasmLoweredZero:
-		getReg32(s, v.Args[0])
+		getValue32(s, v.Args[0])
 		i32Const(s, int32(v.AuxInt))
 		p := s.Prog(wasm.ACall)
 		p.To = obj.Addr{Type: obj.TYPE_MEM, Name: obj.NAME_EXTERN, Sym: gc.WasmZero}
 
 	case ssa.OpWasmLoweredNilCheck:
-		getReg64(s, v.Args[0])
+		getValue64(s, v.Args[0])
 		s.Prog(wasm.AI64Eqz)
 		s.Prog(wasm.AIf)
 		p := s.Prog(wasm.ACALLNORESUME)
@@ -161,14 +165,14 @@ func ssaGenValue(s *gc.SSAGenState, v *ssa.Value) {
 		}
 
 	case ssa.OpWasmLoweredWB:
-		getReg64(s, v.Args[0])
-		getReg64(s, v.Args[1])
+		getValue64(s, v.Args[0])
+		getValue64(s, v.Args[1])
 		p := s.Prog(wasm.ACALLNORESUME) // TODO(neelance): If possible, turn this into a simple wasm.ACall).
 		p.To = obj.Addr{Type: obj.TYPE_MEM, Name: obj.NAME_EXTERN, Sym: v.Aux.(*obj.LSym)}
 
 	case ssa.OpWasmI64Store8, ssa.OpWasmI64Store16, ssa.OpWasmI64Store32, ssa.OpWasmI64Store, ssa.OpWasmF32Store, ssa.OpWasmF64Store:
-		getReg32(s, v.Args[0])
-		getReg64(s, v.Args[1])
+		getValue32(s, v.Args[0])
+		getValue64(s, v.Args[1])
 		if v.Op == ssa.OpWasmF32Store {
 			s.Prog(wasm.AF32DemoteF64)
 		}
@@ -177,7 +181,7 @@ func ssaGenValue(s *gc.SSAGenState, v *ssa.Value) {
 
 	case ssa.OpStoreReg:
 		getReg(s, wasm.REG_SP)
-		getReg64(s, v.Args[0])
+		getValue64(s, v.Args[0])
 		if v.Type.Etype == types.TFLOAT32 {
 			s.Prog(wasm.AF32DemoteF64)
 		}
@@ -188,7 +192,16 @@ func ssaGenValue(s *gc.SSAGenState, v *ssa.Value) {
 		if v.Type.IsMemory() {
 			return
 		}
+		if v.OnWasmStack {
+			s.OnWasmStackSkipped++
+			// If a Value is marked OnWasmStack, we don't generate the value and store it to a register now.
+			// Instead, we delay the generation to when the value is used and then directly generate it on the WebAssembly stack.
+			return
+		}
 		ssaGenValueOnStack(s, v)
+		if s.OnWasmStackSkipped != 0 {
+			panic("wasm: bad stack")
+		}
 		setReg(s, v.Reg())
 	}
 }
@@ -237,22 +250,22 @@ func ssaGenValueOnStack(s *gc.SSAGenState, v *ssa.Value) {
 		}
 
 	case ssa.OpWasmLoweredRound32F:
-		getReg64(s, v.Args[0])
+		getValue64(s, v.Args[0])
 		s.Prog(wasm.AF32DemoteF64)
 		s.Prog(wasm.AF64PromoteF32)
 
 	case ssa.OpWasmLoweredConvert:
-		getReg64(s, v.Args[0])
+		getValue64(s, v.Args[0])
 
 	case ssa.OpWasmSelect:
-		getReg64(s, v.Args[0])
-		getReg64(s, v.Args[1])
-		getReg64(s, v.Args[2])
+		getValue64(s, v.Args[0])
+		getValue64(s, v.Args[1])
+		getValue64(s, v.Args[2])
 		s.Prog(wasm.AI32WrapI64)
 		s.Prog(v.Op.Asm())
 
 	case ssa.OpWasmI64AddConst:
-		getReg64(s, v.Args[0])
+		getValue64(s, v.Args[0])
 		i64Const(s, v.AuxInt)
 		s.Prog(v.Op.Asm())
 
@@ -263,7 +276,7 @@ func ssaGenValueOnStack(s *gc.SSAGenState, v *ssa.Value) {
 		f64Const(s, v.AuxFloat())
 
 	case ssa.OpWasmI64Load8U, ssa.OpWasmI64Load8S, ssa.OpWasmI64Load16U, ssa.OpWasmI64Load16S, ssa.OpWasmI64Load32U, ssa.OpWasmI64Load32S, ssa.OpWasmI64Load, ssa.OpWasmF32Load, ssa.OpWasmF64Load:
-		getReg32(s, v.Args[0])
+		getValue32(s, v.Args[0])
 		p := s.Prog(v.Op.Asm())
 		p.From = obj.Addr{Type: obj.TYPE_CONST, Offset: v.AuxInt}
 		if v.Op == ssa.OpWasmF32Load {
@@ -271,24 +284,24 @@ func ssaGenValueOnStack(s *gc.SSAGenState, v *ssa.Value) {
 		}
 
 	case ssa.OpWasmI64Eqz:
-		getReg64(s, v.Args[0])
+		getValue64(s, v.Args[0])
 		s.Prog(v.Op.Asm())
 		s.Prog(wasm.AI64ExtendUI32)
 
 	case ssa.OpWasmI64Eq, ssa.OpWasmI64Ne, ssa.OpWasmI64LtS, ssa.OpWasmI64LtU, ssa.OpWasmI64GtS, ssa.OpWasmI64GtU, ssa.OpWasmI64LeS, ssa.OpWasmI64LeU, ssa.OpWasmI64GeS, ssa.OpWasmI64GeU, ssa.OpWasmF64Eq, ssa.OpWasmF64Ne, ssa.OpWasmF64Lt, ssa.OpWasmF64Gt, ssa.OpWasmF64Le, ssa.OpWasmF64Ge:
-		getReg64(s, v.Args[0])
-		getReg64(s, v.Args[1])
+		getValue64(s, v.Args[0])
+		getValue64(s, v.Args[1])
 		s.Prog(v.Op.Asm())
 		s.Prog(wasm.AI64ExtendUI32)
 
 	case ssa.OpWasmI64Add, ssa.OpWasmI64Sub, ssa.OpWasmI64Mul, ssa.OpWasmI64DivU, ssa.OpWasmI64RemS, ssa.OpWasmI64RemU, ssa.OpWasmI64And, ssa.OpWasmI64Or, ssa.OpWasmI64Xor, ssa.OpWasmI64Shl, ssa.OpWasmI64ShrS, ssa.OpWasmI64ShrU, ssa.OpWasmF64Add, ssa.OpWasmF64Sub, ssa.OpWasmF64Mul, ssa.OpWasmF64Div:
-		getReg64(s, v.Args[0])
-		getReg64(s, v.Args[1])
+		getValue64(s, v.Args[0])
+		getValue64(s, v.Args[1])
 		s.Prog(v.Op.Asm())
 
 	case ssa.OpWasmI64DivS:
-		getReg64(s, v.Args[0])
-		getReg64(s, v.Args[1])
+		getValue64(s, v.Args[0])
+		getValue64(s, v.Args[1])
 		if v.Type.Size() == 8 {
 			// Division of int64 needs helper function wasmDiv to handle the MinInt64 / -1 case.
 			p := s.Prog(wasm.ACall)
@@ -298,17 +311,17 @@ func ssaGenValueOnStack(s *gc.SSAGenState, v *ssa.Value) {
 		s.Prog(wasm.AI64DivS)
 
 	case ssa.OpWasmI64TruncSF64:
-		getReg64(s, v.Args[0])
+		getValue64(s, v.Args[0])
 		p := s.Prog(wasm.ACall)
 		p.To = obj.Addr{Type: obj.TYPE_MEM, Name: obj.NAME_EXTERN, Sym: gc.WasmTruncS}
 
 	case ssa.OpWasmI64TruncUF64:
-		getReg64(s, v.Args[0])
+		getValue64(s, v.Args[0])
 		p := s.Prog(wasm.ACall)
 		p.To = obj.Addr{Type: obj.TYPE_MEM, Name: obj.NAME_EXTERN, Sym: gc.WasmTruncU}
 
 	case ssa.OpWasmF64Neg, ssa.OpWasmF64ConvertSI64, ssa.OpWasmF64ConvertUI64:
-		getReg64(s, v.Args[0])
+		getValue64(s, v.Args[0])
 		s.Prog(v.Op.Asm())
 
 	case ssa.OpLoadReg:
@@ -319,7 +332,7 @@ func ssaGenValueOnStack(s *gc.SSAGenState, v *ssa.Value) {
 		}
 
 	case ssa.OpCopy:
-		getReg64(s, v.Args[0])
+		getValue64(s, v.Args[0])
 
 	default:
 		v.Fatalf("unexpected op: %s", v.Op)
@@ -327,7 +340,14 @@ func ssaGenValueOnStack(s *gc.SSAGenState, v *ssa.Value) {
 	}
 }
 
-func getReg32(s *gc.SSAGenState, v *ssa.Value) {
+func getValue32(s *gc.SSAGenState, v *ssa.Value) {
+	if v.OnWasmStack {
+		s.OnWasmStackSkipped--
+		ssaGenValueOnStack(s, v)
+		s.Prog(wasm.AI32WrapI64)
+		return
+	}
+
 	reg := v.Reg()
 	getReg(s, reg)
 	if reg != wasm.REG_SP {
@@ -335,7 +355,13 @@ func getReg32(s *gc.SSAGenState, v *ssa.Value) {
 	}
 }
 
-func getReg64(s *gc.SSAGenState, v *ssa.Value) {
+func getValue64(s *gc.SSAGenState, v *ssa.Value) {
+	if v.OnWasmStack {
+		s.OnWasmStackSkipped--
+		ssaGenValueOnStack(s, v)
+		return
+	}
+
 	reg := v.Reg()
 	getReg(s, reg)
 	if reg == wasm.REG_SP {