cmd/compile: rework/reduce partially lived argument spilling

In CL 307909 we generate code that spills pointer-typed argument
registers if it is part of an SSA-able aggregate. The current
code spill the register unconditionally. Sometimes it is
unnecessary, because it is already spilled, or it is never live.

This CL reworks the spill generation. We move it to the end of
compilation, after liveness analysis, so we have information about
if a spill is necessary, and only generate spills for the
necessary ones.

Change-Id: I8d60be9b2c47651aeda14f5e2d1bbd207c134b26
Reviewed-on: https://go-review.googlesource.com/c/go/+/309331
Trust: Cherry Zhang <cherryyz@google.com>
Run-TryBot: Cherry Zhang <cherryyz@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: David Chase <drchase@google.com>
Reviewed-by: Than McIntosh <thanm@google.com>

src/cmd/compile/internal/amd64/galign.go

@@ -24,4 +24,5 @@ func Init(arch *ssagen.ArchInfo) {
 	arch.SSAGenValue = ssaGenValue
 	arch.SSAGenBlock = ssaGenBlock
 	arch.LoadRegResults = loadRegResults
+	arch.SpillArgReg = spillArgReg
 }

src/cmd/compile/internal/amd64/ssa.go

@@ -11,6 +11,7 @@ import (
 	"cmd/compile/internal/base"
 	"cmd/compile/internal/ir"
 	"cmd/compile/internal/logopt"
+	"cmd/compile/internal/objw"
 	"cmd/compile/internal/ssa"
 	"cmd/compile/internal/ssagen"
 	"cmd/compile/internal/types"
@@ -1364,3 +1365,10 @@ func loadRegResults(s *ssagen.State, f *ssa.Func) {
 		}
 	}
 }
+
+func spillArgReg(pp *objw.Progs, p *obj.Prog, f *ssa.Func, t *types.Type, reg int16, n *ir.Name, off int64) *obj.Prog {
+	p = pp.Append(p, storeByType(t), obj.TYPE_REG, reg, 0, obj.TYPE_MEM, 0, n.FrameOffset()+off)
+	p.To.Name = obj.NAME_PARAM
+	p.To.Sym = n.Linksym()
+	return p
+}

src/cmd/compile/internal/liveness/plive.go

@@ -141,6 +141,11 @@ type liveness struct {
 
 	cache progeffectscache
 
+	// partLiveArgs includes input arguments (PPARAM) that may
+	// be partially live. That is, it is considered live because
+	// a part of it is used, but we may not initialize all parts.
+	partLiveArgs map[*ir.Name]bool
+
 	doClobber bool // Whether to clobber dead stack slots in this function.
 }
 
@@ -268,6 +273,12 @@ func (lv *liveness) valueEffects(v *ssa.Value) (int32, liveEffect) {
 		}
 	}
 
+	if n.Class == ir.PPARAM && !n.Addrtaken() && n.Type().Width > int64(types.PtrSize) {
+		// Only aggregate-typed arguments that are not address-taken can be
+		// partially live.
+		lv.partLiveArgs[n] = true
+	}
+
 	var effect liveEffect
 	// Read is a read, obviously.
 	//
@@ -394,6 +405,8 @@ func newliveness(fn *ir.Func, f *ssa.Func, vars []*ir.Name, idx map[*ir.Name]int
 
 	lv.markUnsafePoints()
 
+	lv.partLiveArgs = make(map[*ir.Name]bool)
+
 	lv.enableClobber()
 
 	return lv
@@ -1310,8 +1323,9 @@ func (lv *liveness) emit() (argsSym, liveSym *obj.LSym) {
 // Entry pointer for Compute analysis. Solves for the Compute of
 // pointer variables in the function and emits a runtime data
 // structure read by the garbage collector.
-// Returns a map from GC safe points to their corresponding stack map index.
-func Compute(curfn *ir.Func, f *ssa.Func, stkptrsize int64, pp *objw.Progs) Map {
+// Returns a map from GC safe points to their corresponding stack map index,
+// and a map that contains all input parameters that may be partially live.
+func Compute(curfn *ir.Func, f *ssa.Func, stkptrsize int64, pp *objw.Progs) (Map, map[*ir.Name]bool) {
 	// Construct the global liveness state.
 	vars, idx := getvariables(curfn)
 	lv := newliveness(curfn, f, vars, idx, stkptrsize)
@@ -1373,7 +1387,7 @@ func Compute(curfn *ir.Func, f *ssa.Func, stkptrsize int64, pp *objw.Progs) Map
 		p.To.Sym = x
 	}
 
-	return lv.livenessMap
+	return lv.livenessMap, lv.partLiveArgs
 }
 
 func (lv *liveness) emitStackObjects() *obj.LSym {

src/cmd/compile/internal/ssagen/arch.go

@@ -5,8 +5,10 @@
 package ssagen
 
 import (
+	"cmd/compile/internal/ir"
 	"cmd/compile/internal/objw"
 	"cmd/compile/internal/ssa"
+	"cmd/compile/internal/types"
 	"cmd/internal/obj"
 )
 
@@ -44,4 +46,7 @@ type ArchInfo struct {
 	// into registers. They are already in memory (PPARAMOUT nodes).
 	// Used in open-coded defer return path.
 	LoadRegResults func(s *State, f *ssa.Func)
+
+	// SpillArgReg emits instructions that spill reg to n+off.
+	SpillArgReg func(pp *objw.Progs, p *obj.Prog, f *ssa.Func, t *types.Type, reg int16, n *ir.Name, off int64) *obj.Prog
 }

src/cmd/compile/internal/ssagen/ssa.go

@@ -558,11 +558,6 @@ func buildssa(fn *ir.Func, worker int) *ssa.Func {
 				v := s.newValue0A(ssa.OpArg, n.Type(), n)
 				s.vars[n] = v
 				s.addNamedValue(n, v) // This helps with debugging information, not needed for compilation itself.
-				// TODO(register args) Make liveness more fine-grained to that partial spilling is okay.
-				paramAssignment := ssa.ParamAssignmentForArgName(s.f, n)
-				if len(paramAssignment.Registers) > 1 && n.Type().HasPointers() { // 1 cannot be partially live
-					s.storeParameterRegsToStack(s.f.ABISelf, paramAssignment, n, s.decladdrs[n], true)
-				}
 			} else { // address was taken AND/OR too large for SSA
 				paramAssignment := ssa.ParamAssignmentForArgName(s.f, n)
 				if len(paramAssignment.Registers) > 0 {
@@ -6440,6 +6435,10 @@ type State struct {
 	// liveness analysis.
 	livenessMap liveness.Map
 
+	// partLiveArgs includes arguments that may be partially live, for which we
+	// need to generate instructions that spill the argument registers.
+	partLiveArgs map[*ir.Name]bool
+
 	// lineRunStart records the beginning of the current run of instructions
 	// within a single block sharing the same line number
 	// Used to move statement marks to the beginning of such runs.
@@ -6525,7 +6524,7 @@ func genssa(f *ssa.Func, pp *objw.Progs) {
 
 	e := f.Frontend().(*ssafn)
 
-	s.livenessMap = liveness.Compute(e.curfn, f, e.stkptrsize, pp)
+	s.livenessMap, s.partLiveArgs = liveness.Compute(e.curfn, f, e.stkptrsize, pp)
 
 	openDeferInfo := e.curfn.LSym.Func().OpenCodedDeferInfo
 	if openDeferInfo != nil {
@@ -6867,10 +6866,60 @@ func defframe(s *State, e *ssafn, f *ssa.Func) {
 	pp.Text.To.Val = int32(types.Rnd(f.OwnAux.ArgWidth(), int64(types.RegSize)))
 	pp.Text.To.Offset = frame
 
+	p := pp.Text
+
+	// Insert code to spill argument registers if the named slot may be partially
+	// live. That is, the named slot is considered live by liveness analysis,
+	// (because a part of it is live), but we may not spill all parts into the
+	// slot. This can only happen with aggregate-typed arguments that are SSA-able
+	// and not address-taken (for non-SSA-able or address-taken arguments we always
+	// spill upfront).
+	// TODO(register args) Make liveness more fine-grained to that partial spilling is okay.
+	if objabi.Experiment.RegabiArgs {
+		// First, see if it is already spilled before it may be live. Look for a spill
+		// in the entry block up to the first safepoint.
+		type nameOff struct {
+			n   *ir.Name
+			off int64
+		}
+		partLiveArgsSpilled := make(map[nameOff]bool)
+		for _, v := range f.Entry.Values {
+			if v.Op.IsCall() {
+				break
+			}
+			if v.Op != ssa.OpStoreReg || v.Args[0].Op != ssa.OpArgIntReg {
+				continue
+			}
+			n, off := ssa.AutoVar(v)
+			if n.Class != ir.PPARAM || n.Addrtaken() || !TypeOK(n.Type()) || !s.partLiveArgs[n] {
+				continue
+			}
+			partLiveArgsSpilled[nameOff{n, off}] = true
+		}
+
+		// Then, insert code to spill registers if not already.
+		for _, a := range f.OwnAux.ABIInfo().InParams() {
+			n, ok := a.Name.(*ir.Name)
+			if !ok || n.Addrtaken() || !TypeOK(n.Type()) || !s.partLiveArgs[n] || len(a.Registers) <= 1 {
+				continue
+			}
+			rts, offs := a.RegisterTypesAndOffsets()
+			for i := range a.Registers {
+				if !rts[i].HasPointers() {
+					continue
+				}
+				if partLiveArgsSpilled[nameOff{n, offs[i]}] {
+					continue // already spilled
+				}
+				reg := ssa.ObjRegForAbiReg(a.Registers[i], f.Config)
+				p = Arch.SpillArgReg(pp, p, f, rts[i], reg, n, offs[i])
+			}
+		}
+	}
+
 	// Insert code to zero ambiguously live variables so that the
 	// garbage collector only sees initialized values when it
 	// looks for pointers.
-	p := pp.Text
 	var lo, hi int64
 
 	// Opaque state for backend to use. Current backends use it to