runtime: asynchronous preemption function for x86

This adds asynchronous preemption function for amd64 and 386. These
functions spill and restore all register state that can be used by
user Go code.

For the moment we stub out the other arches.

For #10958, #24543.

Change-Id: I6f93fabe9875f4834922a5712362e79045c00aca
Reviewed-on: https://go-review.googlesource.com/c/go/+/201759
Run-TryBot: Austin Clements <austin@google.com>
Reviewed-by: Cherry Zhang <cherryyz@google.com>

src/cmd/compile/internal/ssa/gen/386Ops.go

@@ -45,6 +45,8 @@ var regNames386 = []string{
 	"X6",
 	"X7",
 
+	// If you add registers, update asyncPreempt in runtime
+
 	// pseudo-registers
 	"SB",
 }

src/cmd/compile/internal/ssa/gen/AMD64Ops.go

@@ -63,6 +63,8 @@ var regNamesAMD64 = []string{
 	"X14",
 	"X15",
 
+	// If you add registers, update asyncPreempt in runtime
+
 	// pseudo-registers
 	"SB",
 }

src/runtime/mkpreempt.go

@@ -0,0 +1,266 @@
+// Copyright 2019 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// +build ignore
+
+// mkpreempt generates the asyncPreempt functions for each
+// architecture.
+package main
+
+import (
+	"flag"
+	"fmt"
+	"io"
+	"log"
+	"os"
+	"strings"
+)
+
+// Copied from cmd/compile/internal/ssa/gen/*Ops.go
+
+var regNames386 = []string{
+	"AX",
+	"CX",
+	"DX",
+	"BX",
+	"SP",
+	"BP",
+	"SI",
+	"DI",
+	"X0",
+	"X1",
+	"X2",
+	"X3",
+	"X4",
+	"X5",
+	"X6",
+	"X7",
+}
+
+var regNamesAMD64 = []string{
+	"AX",
+	"CX",
+	"DX",
+	"BX",
+	"SP",
+	"BP",
+	"SI",
+	"DI",
+	"R8",
+	"R9",
+	"R10",
+	"R11",
+	"R12",
+	"R13",
+	"R14",
+	"R15",
+	"X0",
+	"X1",
+	"X2",
+	"X3",
+	"X4",
+	"X5",
+	"X6",
+	"X7",
+	"X8",
+	"X9",
+	"X10",
+	"X11",
+	"X12",
+	"X13",
+	"X14",
+	"X15",
+}
+
+var out io.Writer
+
+var arches = map[string]func(){
+	"386":     gen386,
+	"amd64":   genAMD64,
+	"arm":     notImplemented,
+	"arm64":   notImplemented,
+	"mips64x": notImplemented,
+	"mipsx":   notImplemented,
+	"ppc64x":  notImplemented,
+	"s390x":   notImplemented,
+	"wasm":    genWasm,
+}
+var beLe = map[string]bool{"mips64x": true, "mipsx": true, "ppc64x": true}
+
+func main() {
+	flag.Parse()
+	if flag.NArg() > 0 {
+		out = os.Stdout
+		for _, arch := range flag.Args() {
+			gen, ok := arches[arch]
+			if !ok {
+				log.Fatalf("unknown arch %s", arch)
+			}
+			header(arch)
+			gen()
+		}
+		return
+	}
+
+	for arch, gen := range arches {
+		f, err := os.Create(fmt.Sprintf("preempt_%s.s", arch))
+		if err != nil {
+			log.Fatal(err)
+		}
+		out = f
+		header(arch)
+		gen()
+		if err := f.Close(); err != nil {
+			log.Fatal(err)
+		}
+	}
+}
+
+func header(arch string) {
+	fmt.Fprintf(out, "// Code generated by mkpreempt.go; DO NOT EDIT.\n\n")
+	if beLe[arch] {
+		base := arch[:len(arch)-1]
+		fmt.Fprintf(out, "// +build %s %sle\n\n", base, base)
+	}
+	fmt.Fprintf(out, "#include \"go_asm.h\"\n")
+	fmt.Fprintf(out, "#include \"textflag.h\"\n\n")
+	fmt.Fprintf(out, "TEXT ·asyncPreempt(SB),NOSPLIT|NOFRAME,$0-0\n")
+}
+
+func p(f string, args ...interface{}) {
+	fmted := fmt.Sprintf(f, args...)
+	fmt.Fprintf(out, "\t%s\n", strings.Replace(fmted, "\n", "\n\t", -1))
+}
+
+type layout struct {
+	stack int
+	regs  []regPos
+}
+
+type regPos struct {
+	pos int
+
+	op  string
+	reg string
+
+	// If this register requires special save and restore, these
+	// give those operations with a %d placeholder for the stack
+	// offset.
+	save, restore string
+}
+
+func (l *layout) add(op, reg string, size int) {
+	l.regs = append(l.regs, regPos{op: op, reg: reg, pos: l.stack})
+	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.stack += size
+}
+
+func (l *layout) save() {
+	for _, reg := range l.regs {
+		if reg.save != "" {
+			p(reg.save, reg.pos)
+		} else {
+			p("%s %s, %d(SP)", reg.op, reg.reg, reg.pos)
+		}
+	}
+}
+
+func (l *layout) restore() {
+	for i := len(l.regs) - 1; i >= 0; i-- {
+		reg := l.regs[i]
+		if reg.restore != "" {
+			p(reg.restore, reg.pos)
+		} else {
+			p("%s %d(SP), %s", reg.op, reg.pos, reg.reg)
+		}
+	}
+}
+
+func gen386() {
+	p("PUSHFL")
+
+	// Save general purpose registers.
+	var l layout
+	for _, reg := range regNames386 {
+		if reg == "SP" || strings.HasPrefix(reg, "X") {
+			continue
+		}
+		l.add("MOVL", reg, 4)
+	}
+
+	// Save the 387 state.
+	l.addSpecial(
+		"FSAVE %d(SP)\nFLDCW runtime·controlWord64(SB)",
+		"FRSTOR %d(SP)",
+		108)
+
+	// Save SSE state only if supported.
+	lSSE := layout{stack: l.stack}
+	for i := 0; i < 8; i++ {
+		lSSE.add("MOVUPS", fmt.Sprintf("X%d", i), 16)
+	}
+
+	p("ADJSP $%d", lSSE.stack)
+	p("NOP SP")
+	l.save()
+	p("CMPB internal∕cpu·X86+const_offsetX86HasSSE2(SB), $1\nJNE nosse")
+	lSSE.save()
+	p("nosse:")
+	p("CALL ·asyncPreempt2(SB)")
+	p("CMPB internal∕cpu·X86+const_offsetX86HasSSE2(SB), $1\nJNE nosse2")
+	lSSE.restore()
+	p("nosse2:")
+	l.restore()
+	p("ADJSP $%d", -lSSE.stack)
+
+	p("POPFL")
+	p("RET")
+}
+
+func genAMD64() {
+	// Assign stack offsets.
+	var l layout
+	for _, reg := range regNamesAMD64 {
+		if reg == "SP" || reg == "BP" {
+			continue
+		}
+		if strings.HasPrefix(reg, "X") {
+			l.add("MOVUPS", reg, 16)
+		} else {
+			l.add("MOVQ", reg, 8)
+		}
+	}
+
+	// TODO: MXCSR register?
+
+	p("PUSHQ BP")
+	p("MOVQ SP, BP")
+	p("// Save flags before clobbering them")
+	p("PUSHFQ")
+	p("// obj doesn't understand ADD/SUB on SP, but does understand ADJSP")
+	p("ADJSP $%d", l.stack)
+	p("// But vet doesn't know ADJSP, so suppress vet stack checking")
+	p("NOP SP")
+	l.save()
+	p("CALL ·asyncPreempt2(SB)")
+	l.restore()
+	p("ADJSP $%d", -l.stack)
+	p("POPFQ")
+	p("POPQ BP")
+	p("RET")
+}
+
+func genWasm() {
+	p("// No async preemption on wasm")
+	p("UNDEF")
+}
+
+func notImplemented() {
+	p("// Not implemented yet")
+	p("JMP ·abort(SB)")
+}

src/runtime/preempt.go

@@ -232,3 +232,18 @@ func resumeG(state suspendGState) {
 func canPreemptM(mp *m) bool {
 	return mp.locks == 0 && mp.mallocing == 0 && mp.preemptoff == "" && mp.p.ptr().status == _Prunning
 }
+
+//go:generate go run mkpreempt.go
+
+// asyncPreempt saves all user registers and calls asyncPreempt2.
+//
+// When stack scanning encounters an asyncPreempt frame, it scans that
+// frame and its parent frame conservatively.
+//
+// asyncPreempt is implemented in assembly.
+func asyncPreempt()
+
+//go:nosplit
+func asyncPreempt2() {
+	// TODO: Enter scheduler
+}

src/runtime/preempt_386.s

@@ -0,0 +1,52 @@
+// Code generated by mkpreempt.go; DO NOT EDIT.
+
+#include "go_asm.h"
+#include "textflag.h"
+
+TEXT ·asyncPreempt(SB),NOSPLIT|NOFRAME,$0-0
+	PUSHFL
+	ADJSP $264
+	NOP SP
+	MOVL AX, 0(SP)
+	MOVL CX, 4(SP)
+	MOVL DX, 8(SP)
+	MOVL BX, 12(SP)
+	MOVL BP, 16(SP)
+	MOVL SI, 20(SP)
+	MOVL DI, 24(SP)
+	FSAVE 28(SP)
+	FLDCW runtime·controlWord64(SB)
+	CMPB internal∕cpu·X86+const_offsetX86HasSSE2(SB), $1
+	JNE nosse
+	MOVUPS X0, 136(SP)
+	MOVUPS X1, 152(SP)
+	MOVUPS X2, 168(SP)
+	MOVUPS X3, 184(SP)
+	MOVUPS X4, 200(SP)
+	MOVUPS X5, 216(SP)
+	MOVUPS X6, 232(SP)
+	MOVUPS X7, 248(SP)
+	nosse:
+	CALL ·asyncPreempt2(SB)
+	CMPB internal∕cpu·X86+const_offsetX86HasSSE2(SB), $1
+	JNE nosse2
+	MOVUPS 248(SP), X7
+	MOVUPS 232(SP), X6
+	MOVUPS 216(SP), X5
+	MOVUPS 200(SP), X4
+	MOVUPS 184(SP), X3
+	MOVUPS 168(SP), X2
+	MOVUPS 152(SP), X1
+	MOVUPS 136(SP), X0
+	nosse2:
+	FRSTOR 28(SP)
+	MOVL 24(SP), DI
+	MOVL 20(SP), SI
+	MOVL 16(SP), BP
+	MOVL 12(SP), BX
+	MOVL 8(SP), DX
+	MOVL 4(SP), CX
+	MOVL 0(SP), AX
+	ADJSP $-264
+	POPFL
+	RET