runtime: save scalar registers off stack in amd64 async preemption

Asynchronous preemption must save all registers that could be in use
by Go code. Currently, it saves all of these to the goroutine stack.
As a result, the stack frame requirements of asynchronous preemption
can be rather high. On amd64, this requires 368 bytes of stack space,
most of which is the XMM registers. Several RISC architectures are
around 0.5 KiB.

As we add support for SIMD instructions, this is going to become a
problem. The AVX-512 register state is 2.5 KiB. This well exceeds the
nosplit limit, and even if it didn't, could constrain when we can
asynchronously preempt goroutines on small stacks.

This CL fixes this by moving pure scalar state stored in non-GP
registers off the stack and into an allocated "extended register
state" object. To reduce space overhead, we only allocate these
objects as needed. While in the theoretical limit, every G could need
this register state, in practice very few do at a time.

However, we can't allocate when we're in the middle of saving the
register state during an asynchronous preemption, so we reserve
scratch space on every P to temporarily store the register state,
which can then be copied out to an allocated state object later by Go
code.

This commit only implements this for amd64, since that's where we're
about to add much more vector state, but it lays the groundwork for
doing this on any architecture that could benefit.

This is a cherry-pick of CL 680898 plus bug fix CL 684836 from the
dev.simd branch.

Change-Id: I123a95e21c11d5c10942d70e27f84d2d99bbf735
Reviewed-on: https://go-review.googlesource.com/c/go/+/669195
Auto-Submit: Austin Clements <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>
Reviewed-by: Cherry Mui <[email protected]>

Author: aclements

src/runtime/export_test.go

@@ -554,6 +554,8 @@ type G = g
 
 type Sudog = sudog
 
+type XRegPerG = xRegPerG
+
 func Getg() *G {
 	return getg()
 }

src/runtime/lockrank.go

@@ -821,6 +821,8 @@ func (h *mheap) init() {
 	}
 
 	h.pages.init(&h.lock, &memstats.gcMiscSys, false)
+
+	xRegInitAlloc()
 }
 
 // reclaim sweeps and reclaims at least npage pages into the heap.

src/runtime/mheap.go

@@ -193,6 +193,9 @@ defer,
 # Below WB is the write barrier implementation.
 < wbufSpans;
 
+# xRegState allocator
+sched < xRegAlloc;
+
 # Span allocator
 stackLarge,
   stackpool,
@@ -205,7 +208,8 @@ stackLarge,
 # an mspanSpecial lock, and they're part of the malloc implementation.
 # Pinner bits might be freed by the span allocator.
 mheap, mspanSpecial < mheapSpecial;
-mheap, mheapSpecial < globalAlloc;
+# Fixallocs
+mheap, mheapSpecial, xRegAlloc < globalAlloc;
 
 # Execution tracer events (with a P)
 hchan,

src/runtime/mklockrank.go

@@ -9,8 +9,10 @@
 package main
 
 import (
+	"bytes"
 	"flag"
 	"fmt"
+	"go/format"
 	"io"
 	"log"
 	"os"
@@ -122,14 +124,19 @@ type gen struct {
 	goarch string
 }
 
-func (g *gen) asmHeader() {
+func (g *gen) commonHeader() {
 	fmt.Fprintf(g.w, "// Code generated by mkpreempt.go; DO NOT EDIT.\n\n")
 	if beLe[g.goarch] {
 		base := g.goarch[:len(g.goarch)-1]
 		fmt.Fprintf(g.w, "//go:build %s || %sle\n\n", base, base)
 	}
+}
+
+func (g *gen) asmHeader() {
+	g.commonHeader()
 	fmt.Fprintf(g.w, "#include \"go_asm.h\"\n")
 	if g.goarch == "amd64" {
+		fmt.Fprintf(g.w, "#include \"go_tls.h\"\n")
 		fmt.Fprintf(g.w, "#include \"asm_amd64.h\"\n")
 	}
 	fmt.Fprintf(g.w, "#include \"textflag.h\"\n\n")
@@ -145,14 +152,51 @@ func (g *gen) label(l string) {
 	fmt.Fprintf(g.w, "%s\n", l)
 }
 
+// writeXRegs writes an architecture xregs file.
+func writeXRegs(arch string, l *layout) {
+	var code bytes.Buffer
+	g := gen{&code, arch}
+	g.commonHeader()
+	fmt.Fprintf(g.w, `
+package runtime
+
+type xRegs struct {
+`)
+	pos := 0
+	for _, reg := range l.regs {
+		if reg.pos != pos {
+			log.Fatalf("padding not implemented")
+		}
+		typ := fmt.Sprintf("[%d]byte", reg.size)
+		switch {
+		case reg.size == 4 && reg.pos%4 == 0:
+			typ = "uint32"
+		case reg.size == 8 && reg.pos%8 == 0:
+			typ = "uint64"
+		}
+		fmt.Fprintf(g.w, "\t%s %s\n", reg.reg, typ)
+		pos += reg.size
+	}
+	fmt.Fprintf(g.w, "}\n")
+
+	path := fmt.Sprintf("preempt_%s.go", arch)
+	b, err := format.Source(code.Bytes())
+	if err != nil {
+		log.Fatalf("formatting %s: %s", path, err)
+	}
+	if err := os.WriteFile(path, b, 0666); err != nil {
+		log.Fatal(err)
+	}
+}
+
 type layout struct {
 	stack int
 	regs  []regPos
 	sp    string // stack pointer register
 }
 
 type regPos struct {
-	pos int
+	pos, size int
 
 	saveOp    string
 	restoreOp string
@@ -165,17 +209,17 @@ type regPos struct {
 }
 
 func (l *layout) add(op, reg string, size int) {
-	l.regs = append(l.regs, regPos{saveOp: op, restoreOp: op, reg: reg, pos: l.stack})
+	l.regs = append(l.regs, regPos{saveOp: op, restoreOp: op, reg: reg, pos: l.stack, size: size})
 	l.stack += size
 }
 
 func (l *layout) add2(sop, rop, reg string, size int) {
-	l.regs = append(l.regs, regPos{saveOp: sop, restoreOp: rop, reg: reg, pos: l.stack})
+	l.regs = append(l.regs, regPos{saveOp: sop, restoreOp: rop, reg: reg, pos: l.stack, size: size})
 	l.stack += size
 }
 
 func (l *layout) addSpecial(save, restore string, size int) {
-	l.regs = append(l.regs, regPos{save: save, restore: restore, pos: l.stack})
+	l.regs = append(l.regs, regPos{save: save, restore: restore, pos: l.stack, size: size})
 	l.stack += size
 }
 
@@ -239,6 +283,8 @@ func gen386(g *gen) {
 }
 
 func genAMD64(g *gen) {
+	const xReg = "AX" // *xRegState
+
 	p := g.p
 
 	// Assign stack offsets.
@@ -251,12 +297,13 @@ func genAMD64(g *gen) {
 			l.add("MOVQ", reg, 8)
 		}
 	}
-	lSSE := layout{stack: l.stack, sp: "SP"}
+	lXRegs := layout{sp: xReg} // Non-GP registers
 	for _, reg := range regNamesAMD64 {
 		if strings.HasPrefix(reg, "X") {
-			lSSE.add("MOVUPS", reg, 16)
+			lXRegs.add("MOVUPS", reg, 16)
 		}
 	}
+	writeXRegs(g.goarch, &lXRegs)
 
 	// TODO: MXCSR register?
 
@@ -265,17 +312,40 @@ func genAMD64(g *gen) {
 	p("// Save flags before clobbering them")
 	p("PUSHFQ")
 	p("// obj doesn't understand ADD/SUB on SP, but does understand ADJSP")
-	p("ADJSP $%d", lSSE.stack)
+	p("ADJSP $%d", l.stack)
 	p("// But vet doesn't know ADJSP, so suppress vet stack checking")
 	p("NOP SP")
 
+	p("// Save GPs")
 	l.save(g)
 
-	lSSE.save(g)
+	// In general, the limitations on asynchronous preemption mean we only
+	// preempt in ABIInternal code. However, there's at least one exception to
+	// this: when we're in an open-coded transition between an ABIInternal
+	// function and an ABI0 call. We could more carefully arrange unsafe points
+	// to avoid ever landing in ABI0, but it's easy to just make this code not
+	// sensitive to the ABI we're preempting. The CALL to asyncPreempt2 will
+	// ensure we're in ABIInternal register state.
+	p("// Save extended register state to p.xRegs.scratch")
+	p("// Don't make assumptions about ABI register state. See mkpreempt.go")
+	p("get_tls(CX)")
+	p("MOVQ g(CX), R14")
+	p("MOVQ g_m(R14), %s", xReg)
+	p("MOVQ m_p(%s), %s", xReg, xReg)
+	p("LEAQ (p_xRegs+xRegPerP_scratch)(%s), %s", xReg, xReg)
+	lXRegs.save(g)
+
 	p("CALL ·asyncPreempt2(SB)")
-	lSSE.restore(g)
+
+	p("// Restore non-GPs from *p.xRegs.cache")
+	p("MOVQ g_m(R14), %s", xReg)
+	p("MOVQ m_p(%s), %s", xReg, xReg)
+	p("MOVQ (p_xRegs+xRegPerP_cache)(%s), %s", xReg, xReg)
+	lXRegs.restore(g)
+
+	p("// Restore GPs")
 	l.restore(g)
-	p("ADJSP $%d", -lSSE.stack)
+	p("ADJSP $%d", -l.stack)
 	p("POPFQ")
 	p("POPQ BP")
 	p("RET")

src/runtime/mkpreempt.go

@@ -292,21 +292,52 @@ func canPreemptM(mp *m) bool {
 
 // asyncPreempt saves all user registers and calls asyncPreempt2.
 //
-// When stack scanning encounters an asyncPreempt frame, it scans that
+// It saves GP registers (anything that might contain a pointer) to the G stack.
+// Hence, when stack scanning encounters an asyncPreempt frame, it scans that
 // frame and its parent frame conservatively.
 //
+// On some platforms, it saves large additional scalar-only register state such
+// as vector registers to an "extended register state" on the P.
+//
 // asyncPreempt is implemented in assembly.
 func asyncPreempt()
 
+// asyncPreempt2 is the Go continuation of asyncPreempt.
+//
+// It must be deeply nosplit because there's untyped data on the stack from
+// asyncPreempt.
+//
+// It must not have any write barriers because we need to limit the amount of
+// stack it uses.
+//
 //go:nosplit
+//go:nowritebarrierrec
 func asyncPreempt2() {
+	// We can't grow the stack with untyped data from asyncPreempt, so switch to
+	// the system stack right away.
+	mcall(func(gp *g) {
+		gp.asyncSafePoint = true
+
+		// Move the extended register state from the P to the G. We do this now that
+		// we're on the system stack to avoid stack splits.
+		xRegSave(gp)
+
+		if gp.preemptStop {
+			preemptPark(gp)
+		} else {
+			gopreempt_m(gp)
+		}
+		// The above functions never return.
+	})
+
+	// Do not grow the stack below here!
+
 	gp := getg()
-	gp.asyncSafePoint = true
-	if gp.preemptStop {
-		mcall(preemptPark)
-	} else {
-		mcall(gopreempt_m)
-	}
+
+	// Put the extended register state back on the M so resumption can find it.
+	// We can't do this in asyncPreemptM because the park calls never return.
+	xRegRestore(gp)
+
 	gp.asyncSafePoint = false
 }
 
@@ -319,19 +350,13 @@ func init() {
 	total := funcMaxSPDelta(f)
 	f = findfunc(abi.FuncPCABIInternal(asyncPreempt2))
 	total += funcMaxSPDelta(f)
+	f = findfunc(abi.FuncPCABIInternal(xRegRestore))
+	total += funcMaxSPDelta(f)
 	// Add some overhead for return PCs, etc.
 	asyncPreemptStack = uintptr(total) + 8*goarch.PtrSize
 	if asyncPreemptStack > stackNosplit {
-		// We need more than the nosplit limit. This isn't
-		// unsafe, but it may limit asynchronous preemption.
-		//
-		// This may be a problem if we start using more
-		// registers. In that case, we should store registers
-		// in a context object. If we pre-allocate one per P,
-		// asyncPreempt can spill just a few registers to the
-		// stack, then grab its context object and spill into
-		// it. When it enters the runtime, it would allocate a
-		// new context for the P.
+		// We need more than the nosplit limit. This isn't unsafe, but it may
+		// limit asynchronous preemption. Consider moving state into xRegState.
 		print("runtime: asyncPreemptStack=", asyncPreemptStack, "\n")
 		throw("async stack too large")
 	}